Backend-aware Like types, BF16Like, D/K dims, tox matrix trim

- Thread `asarray` parameter through _ArrayLikeChecker so JAX/PyTorch Like
  types use jnp.asarray / torch.as_tensor on the slow path, accepting
  __jax_array__ and __torch_function__ protocol objects
- Add exact-match shortcut in _check_dtype for non-numpy dtypes (bfloat16)
- Add BF16Like to jax.py and torch.py (was missing from Like type set)
- Add D (embedding dim) and K (num heads) dimension symbols for transformers
- Export __version__ from shapix.__init__
- Trim tox matrix from ~252 to ~35 envs (oldest+newest strategy)
- Lower beartype minimum to >=0.20 to match tox matrix
- Defensive error in _shape.py fallthrough instead of silent pass
- Document `from __future__ import annotations` compatibility in README

Author: acecchini

README.md

@@ -110,6 +110,8 @@ Bind to a size on first occurrence and enforce consistency on subsequent ones.
 | `N` | Batch size, count |
 | `B` | Batch |
 | `C` | Channels |
+| `D` | Embedding dimension |
+| `K` | Number of heads |
 | `H` | Height |
 | `W` | Width |
 | `L` | Sequence length |
@@ -299,15 +301,15 @@ from shapix.numpy import F32, I64, Shaped  # and many more
 from shapix.jax import F32, BF16
 ```
 
-Same type names as NumPy, plus `BF16`. Base type is `jax.Array`. Also exports `Like` types and `Tree`.
+Same type names as NumPy, plus `BF16` and `BF16Like`. Base type is `jax.Array`. Also exports `Like` types and `Tree`.
 
 ### PyTorch
 
 ```python
 from shapix.torch import F32, BF16
 ```
 
-Same type names as NumPy, plus `BF16`. Base type is `torch.Tensor`. Also exports `Like` types.
+Same type names as NumPy, plus `BF16` and `BF16Like`. Base type is `torch.Tensor`. Also exports `Like` types.
 
 ### Endianness variants
 
@@ -564,6 +566,29 @@ def train(params: Tree[F32[N], Params], state: Tree[I64[N], State]): ...
 
 ## Advanced usage
 
+### `from __future__ import annotations` (PEP 563)
+
+Shapix is fully compatible with `from __future__ import annotations`. The library itself uses it in every source file.
+
+The one rule: **every dimension symbol used in an annotation must be imported in the module scope.** This is true with or without the future import — the difference is only the error you get if you forget:
+
+```python
+from __future__ import annotations
+from shapix import C           # B is NOT imported
+from shapix.numpy import F32
+
+@beartype
+def f(x: F32[~B, C]): ...     # BeartypeDecorHintForwardRefException — B is not in scope
+```
+
+Fix: import `B`:
+
+```python
+from shapix import B, C
+```
+
+This applies to all dimension symbols — named (`N`, `B`), custom (`Vocab = Dimension("Vocab")`), and any symbol used with operators (`~B`, `+N`). The operators (`~`, `+`) are evaluated on the imported object, so the base symbol must be available.
+
 ### Package-wide instrumentation with `beartype.claw`
 
 Instead of decorating each function with `@beartype`, you can instrument an entire package:
@@ -753,7 +778,7 @@ def f(x: F32[~B, C]) -> F32[~B, C]:  # type: ignore[reportInvalidTypeForm]
 
 ### Dimension symbols (`shapix`)
 
-`N`, `B`, `C`, `H`, `W`, `L`, `P` — named dimensions
+`N`, `B`, `C`, `D`, `K`, `H`, `W`, `L`, `P` — named dimensions
 `__` — anonymous dimension
 `Scalar` — scalar (zero dimensions)
 `T`, `S` — tree structure symbols
@@ -775,7 +800,7 @@ def f(x: F32[~B, C]) -> F32[~B, C]:  # type: ignore[reportInvalidTypeForm]
 
 ### JAX/PyTorch (`shapix.jax`, `shapix.torch`)
 
-Same array types as NumPy, plus `BF16`. Both export `Like` types, `ScalarLike` types (re-exported from numpy), and `make_scalar_like_type`. JAX also exports `Tree`.
+Same array types as NumPy, plus `BF16` and `BF16Like`. Both export `Like` types, `ScalarLike` types (re-exported from numpy), and `make_scalar_like_type`. JAX also exports `Tree`.
 
 ### Factories (`shapix`)
 

pyproject.toml

@@ -6,7 +6,7 @@ readme = "README.md"
 license = "MIT"
 authors = [{ name = "acecchini", email = "ale.cecchini.valette@gmail.com" }]
 requires-python = ">=3.12"
-dependencies = ["beartype>=0.22.9"]
+dependencies = ["beartype>=0.20"]
 keywords = [
   "type-checking",
   "shape",

src/shapix/__init__.py

@@ -20,7 +20,7 @@ def conv(x: F32[N, C, H, W]) -> F32[N, C, H, W]: ...
 Exports
 -------
 Dimension symbols
-    ``B``, ``N``, ``P``, ``L``, ``C``, ``H``, ``W`` — named dimensions.
+    ``B``, ``N``, ``P``, ``L``, ``C``, ``D``, ``K``, ``H``, ``W`` — named dimensions.
     ``__`` — anonymous (match any single dim, no binding).
     ``Scalar`` — scalar (no dimensions).
 
@@ -55,6 +55,10 @@ def conv(x: F32[N, C, H, W]) -> F32[N, C, H, W]: ...
     ``is_bearable()`` checks.
 """
 
+from importlib.metadata import version
+
+__version__ = version("shapix")
+
 from ._array_types import make_array_like_type as make_array_like_type
 from ._array_types import make_array_type as make_array_type
 from ._decorator import check as check
@@ -65,8 +69,10 @@ def conv(x: F32[N, C, H, W]) -> F32[N, C, H, W]: ...
 from ._dimensions import __ as __
 from ._dimensions import B as B
 from ._dimensions import C as C
+from ._dimensions import D as D
 from ._dimensions import Dimension as Dimension
 from ._dimensions import H as H
+from ._dimensions import K as K
 from ._dimensions import L as L
 from ._dimensions import N as N
 from ._dimensions import P as P

src/shapix/_array_types.py

@@ -174,13 +174,25 @@ class _ArrayLikeChecker:
 
   Handles arrays (objects with ``.shape`` + ``.dtype``), scalars, nested
   sequences, ``__array__`` protocol objects, and buffer protocol objects
-  by converting to a numpy array and checking dtype + shape.
+  by converting to an array and checking dtype + shape.
 
   The *casting* parameter controls dtype strictness using numpy casting rules:
   ``no < equiv < safe < same_kind < unsafe``.
+
+  The optional *asarray* callable allows backends to provide their own
+  conversion function (e.g. ``jnp.asarray``, ``torch.as_tensor``) so that
+  objects implementing backend-specific protocols (``__jax_array__``,
+  ``__torch_function__``) are accepted.
   """
 
-  __slots__ = ("_dtype_spec", "_shape_spec", "_casting", "_repr", "_fail_obj")
+  __slots__ = (
+    "_dtype_spec",
+    "_shape_spec",
+    "_casting",
+    "_asarray",
+    "_repr",
+    "_fail_obj",
+  )
 
   def __init__(
     self,
@@ -189,10 +201,12 @@ def __init__(
     *,
     casting: str,
     name: str,
+    asarray: object | None = None,
   ) -> None:
     self._dtype_spec = dtype_spec
     self._shape_spec = shape_spec
     self._casting = casting
+    self._asarray = asarray
     self._fail_obj: object | None = None
 
     dims = ", ".join(repr(d) for d in shape_spec)
@@ -215,20 +229,34 @@ def __call__(self, obj: object) -> bool:
         self._fail_obj = obj
       return result
 
-    # Slow path: convert scalar / sequence / __array__ / buffer to numpy array
-    try:
-      import numpy as np
-
-      arr = np.asarray(obj)
-    except (TypeError, ValueError):
+    # Slow path: convert scalar / sequence / protocol object to array.
+    # Try the backend-specific converter first (jnp.asarray, torch.as_tensor),
+    # then fall back to np.asarray for scalars and nested sequences.
+    arr = self._convert(obj)
+    if arr is None:
       self._fail_obj = obj
       return False
 
-    result = self._check(arr, tuple(arr.shape), memo)
+    result = self._check(arr, tuple(arr.shape), memo)  # type: ignore[union-attr]
     if not result:
       self._fail_obj = obj
     return result
 
+  def _convert(self, obj: object) -> object | None:
+    """Convert *obj* to an array with ``.shape`` and ``.dtype``, or None."""
+    if self._asarray is not None:
+      try:
+        return self._asarray(obj)  # type: ignore[operator]
+      except Exception:  # noqa: BLE001
+        pass  # fall through to numpy
+
+    try:
+      import numpy as np
+
+      return np.asarray(obj)
+    except (TypeError, ValueError):
+      return None
+
   def _check(self, obj: object, shape: tuple[int, ...], memo: ShapeMemo) -> bool:
     """Validate dtype (with casting rules) then shape (with memo)."""
     if not self._check_dtype(obj):
@@ -254,6 +282,11 @@ def _check_dtype(self, obj: object) -> bool:
     if "*" in self._dtype_spec.allowed:
       return self._dtype_spec._check_byteorder(obj)
 
+    # Exact string match always passes (handles non-numpy dtypes like bfloat16
+    # where np.can_cast would raise TypeError for an unknown dtype string).
+    if source in self._dtype_spec.allowed:
+      return self._dtype_spec._check_byteorder(obj)
+
     import numpy as np
 
     for target in self._dtype_spec.allowed:
@@ -281,11 +314,19 @@ class _ArrayLikeFactory:
   sequences, arrays) with dtype casting awareness.
   """
 
-  __slots__ = ("_dtype_spec", "_casting", "__name__")
+  __slots__ = ("_dtype_spec", "_casting", "_asarray", "__name__")
 
-  def __init__(self, dtype_spec: DtypeSpec, *, casting: str, name: str) -> None:
+  def __init__(
+    self,
+    dtype_spec: DtypeSpec,
+    *,
+    casting: str,
+    name: str,
+    asarray: object | None = None,
+  ) -> None:
     self._dtype_spec = dtype_spec
     self._casting = casting
+    self._asarray = asarray
     self.__name__ = name
 
   def __getitem__(self, dims: object) -> type:
@@ -294,7 +335,11 @@ def __getitem__(self, dims: object) -> type:
 
     shape_spec = _to_shape_spec(dims)
     checker = _ArrayLikeChecker(
-      self._dtype_spec, shape_spec, casting=self._casting, name=self.__name__
+      self._dtype_spec,
+      shape_spec,
+      casting=self._casting,
+      name=self.__name__,
+      asarray=self._asarray,
     )
     return Annotated[object, Is[checker]]  # type: ignore[return-value]
 
@@ -303,7 +348,11 @@ def __repr__(self) -> str:
 
 
 def make_array_like_type(
-  dtype_spec: DtypeSpec, *, casting: str = "same_kind", name: str = "ArrayLike"
+  dtype_spec: DtypeSpec,
+  *,
+  casting: str = "same_kind",
+  name: str = "ArrayLike",
+  asarray: object | None = None,
 ) -> _ArrayLikeFactory:
   """Create a subscriptable array-like type factory.
 
@@ -317,6 +366,11 @@ def make_array_like_type(
       compatibility is checked.  Default ``"same_kind"``.
   name:
       Human-readable name for error messages.
+  asarray:
+      Optional callable ``(obj) -> array`` for backend-specific conversion.
+      When provided, it is tried before ``np.asarray`` on the slow path
+      (objects without ``.shape`` / ``.dtype``).  Use this to support
+      protocols like ``__jax_array__`` or ``__torch_function__``.
 
   Returns
   -------
@@ -334,7 +388,7 @@ def make_array_like_type(
   if casting not in _VALID_CASTINGS:
     msg = f"Invalid casting {casting!r}, must be one of {sorted(_VALID_CASTINGS)}"
     raise ValueError(msg)
-  return _ArrayLikeFactory(dtype_spec, casting=casting, name=name)
+  return _ArrayLikeFactory(dtype_spec, casting=casting, name=name, asarray=asarray)
 
 
 # ---------------------------------------------------------------------------

src/shapix/_dimensions.py

@@ -51,6 +51,8 @@ def flatten(x: F32[N, C]) -> F32[N * C]: ...
   "P",
   "L",
   "C",
+  "D",
+  "K",
   "H",
   "W",
   # Anonymous
@@ -195,6 +197,8 @@ def _dim_spec(self) -> DimSpec | None:
   P: Dimension
   L: Dimension
   C: Dimension
+  D: Dimension
+  K: Dimension
   H: Dimension
   W: Dimension
   __: Dimension
@@ -206,6 +210,8 @@ def _dim_spec(self) -> DimSpec | None:
   P = Dimension("P")
   L = Dimension("L")
   C = Dimension("C")
+  D = Dimension("D")
+  K = Dimension("K")
   H = Dimension("H")
   W = Dimension("W")
 

src/shapix/_shape.py

@@ -234,7 +234,7 @@ def _check_one(dim: DimSpec, size: int, memo: ShapeMemo) -> str:
       return f"dimension '{dim.expr}' evaluated to {expected} but got {size}"
     return ""
 
-  return ""
+  return f"internal error: unrecognized dim spec {dim!r}"
 
 
 def _check_variadic(dim: VariadicDim, shape: tuple[int, ...], memo: ShapeMemo) -> str: