refactor: unify operator selection behind generic dispatch

- add astrai/extension/dispatch.py: per-family decision tables over composable Specs with explicit-strict / implicit-loose resolution, ASTR_OPS env overrides, profile presets, and explain traces
- make the axis schema family-owned: register_family takes an axes extractor that snapshots whatever decision axes that family needs from the call, and the core only supplies the axis() predicate vocabulary plus a tensor_axes helper
- drop the central CallContext dataclass; resolve and explain take the raw call arguments, so unregistered handles are probed through supports_call on the same args
- migrate attention and rotary onto family-owned axes with behavior-preserving specs and spec-vs-supports_call mirror tests
- replace the non-ASCII member-of glyph in spec descriptions with plain ASCII " in "
This commit is contained in:
2026-09-01 16:56:57 +08:00
parent aabf366633
commit b4d702cd14
5 changed files with 918 additions and 94 deletions
+30
View File
@@ -27,6 +27,22 @@ from astrai.extension.backend import (
attn_backend,
get_backend,
)
from astrai.extension.dispatch import (
Axes,
ExplicitSelectionError,
ImplRecord,
Resolution,
Spec,
axis,
explain,
explain_plan,
op_backend,
register_env_alias,
register_family,
resolve,
resolve_plan,
tensor_axes,
)
from astrai.extension.loader import KERNEL_NAMES, is_available
from astrai.extension.ops import (
TensorLayout,
@@ -52,4 +68,18 @@ __all__ = [
"is_available",
"KERNEL_NAMES",
"apply_rotary_emb",
"Axes",
"ExplicitSelectionError",
"ImplRecord",
"Resolution",
"Spec",
"axis",
"explain",
"explain_plan",
"op_backend",
"register_env_alias",
"register_family",
"resolve",
"resolve_plan",
"tensor_axes",
]