Add symbolic PyTree annotations, expand tox/CI with jax 0.7-0.9 and optree/beartype compat

Replace string-based PyTree structure specs ("T", "S T", "T ...") with
symbolic Structure class and pre-exported T/S symbols. Add composite
(S[T]), prefix (T[...]), and suffix (..., T) operators. Include optree
as optional dependency with full version compat matrix (0.14-0.19),
beartype compat (0.20-0.22), and jax 0.5-0.9.

Author: acecchini

.github/workflows/ci.yml

@@ -47,24 +47,71 @@ jobs:
       fail-fast: false
       matrix:
         python: ["3.12", "3.13", "3.14"]
-        suffix: ["", "-torch", "-jax", "-all"]
-    name: test (py${{ matrix.python }}${{ matrix.suffix }})
+    name: test (py${{ matrix.python }})
     steps:
       - uses: actions/checkout@v6
       - uses: astral-sh/setup-uv@v7
-      - run: uv run tox -e "py${{ matrix.python }}${{ matrix.suffix }}"
+      - run: uv run tox -e "py${{ matrix.python }}"
         env:
           TOX_PYTHON: python${{ matrix.python }}
 
+  numpy-compat:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        version: ["numpy22", "numpy23", "numpy24"]
+    name: ${{ matrix.version }}
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+      - run: uv run tox -e "${{ matrix.version }}"
+
+  jax-compat:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        version: ["jax05", "jax06", "jax07", "jax08", "jax09"]
+    name: ${{ matrix.version }}
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+      - run: uv run tox -e "${{ matrix.version }}"
+
+  torch-compat:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        version: ["torch26", "torch27", "torch28", "torch29", "torch210"]
+    name: ${{ matrix.version }}
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+      - run: uv run tox -e "${{ matrix.version }}"
+
   beartype-compat:
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
         version: ["bt020", "bt021", "bt022"]
-        suffix: ["", "-torch", "-jax"]
-    name: beartype (${{ matrix.version }}${{ matrix.suffix }})
+        suffix: ["", "-jax", "-torch"]
+    name: ${{ matrix.version }}${{ matrix.suffix }}
     steps:
       - uses: actions/checkout@v6
       - uses: astral-sh/setup-uv@v7
       - run: uv run tox -e "${{ matrix.version }}${{ matrix.suffix }}"
+
+  optree-compat:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        version: ["optree014", "optree015", "optree016", "optree017", "optree018", "optree019"]
+    name: ${{ matrix.version }}
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+      - run: uv run tox -e "${{ matrix.version }}"

README.md

@@ -30,11 +30,13 @@ Shapix turns array shape annotations into **Python objects** that beartype valid
 pip install shapix
 ```
 
-Optional backends:
+Shapix has one dependency: [beartype](https://github.com/beartype/beartype). Install your preferred array framework separately:
 
 ```bash
-pip install shapix[jax]    # JAX support
-pip install shapix[torch]  # PyTorch support
+pip install shapix numpy          # NumPy
+pip install shapix torch          # PyTorch
+pip install shapix jax            # JAX
+pip install shapix numpy optree   # NumPy + PyTree support
 ```
 
 ## Quick start
@@ -101,7 +103,7 @@ Bind to a size on first occurrence and enforce consistency on subsequent ones.
 | `H` | Height |
 | `W` | Width |
 | `L` | Sequence length |
-| `T` | Time / tree |
+| `T` | Tree structure |
 | `P` | Points / parameters |
 
 ```python
@@ -322,6 +324,70 @@ BF16_OR_F32 = DtypeSpec("BF16orF32", frozenset({"bfloat16", "float32"}))
 MixedArray = make_array_type(np.ndarray, BF16_OR_F32)
 ```
 
+## PyTree annotations
+
+PyTree annotations validate all leaves in a nested structure (dicts, lists, tuples, namedtuples). Requires `optree` (`pip install optree`).
+
+### Basic leaf checking
+
+```python
+from shapix import PyTree, T, S, N, C
+from shapix.numpy import F32
+
+@beartype
+def process(data: PyTree[F32[N, C]]) -> PyTree[F32[N, C]]:
+    ...
+
+# All leaves must be F32 arrays with consistent N and C
+process({"params": np.ones((3, 4), dtype=np.float32),
+         "state": np.ones((3, 4), dtype=np.float32)})
+```
+
+### Structure binding
+
+Named structure symbols (`T`, `S`) enforce that multiple arguments share identical tree shapes:
+
+```python
+@beartype
+def add_trees(x: PyTree[F32[N], T], y: PyTree[F32[N], T]) -> PyTree[F32[N]]:
+    ...
+
+add_trees({"a": x1, "b": x2}, {"a": y1, "b": y2})  # OK — same structure
+add_trees({"a": x1}, [y1, y2])                       # Raises — different structure
+```
+
+### Composite structures
+
+Use `S[T]` for nested structure matching — outer structure S, inner structure T:
+
+```python
+def f(x: PyTree[int, T], y: PyTree[int, S], z: PyTree[int, S[T]]): ...
+```
+
+### Prefix and suffix wildcards
+
+```python
+# T[...] — top-level structure matches T, subtrees are arbitrary
+def f(x: PyTree[F32[N], T[...]], y: PyTree[F32[N], T[...]]): ...
+
+# ..., T — full structure matches T
+def f(x: PyTree[F32[N], ..., T], y: PyTree[F32[N], ..., T]): ...
+```
+
+### Custom structure symbols
+
+Create your own with `Structure`:
+
+```python
+from shapix import Structure
+
+Params = Structure("Params")
+State = Structure("State")
+
+@beartype
+def train(params: PyTree[F32[N], Params], state: PyTree[I64[N], State]): ...
+```
+
 ## Advanced usage
 
 ### Package-wide instrumentation with `beartype.claw`
@@ -340,28 +406,29 @@ shapix_this_package()
 
 Every function in your package that uses shapix type annotations will be checked automatically.
 
-### Explicit memo management with `@shapix.check`
+### How cross-argument checking works (the memo)
 
-The frame-based memo works automatically with `@beartype` in virtually all cases. For exotic call-stack scenarios, or to combine memo management with custom `BeartypeConf`, use the explicit decorator:
+A **dimension memo** maps dimension names to sizes (e.g., `N→4`, `C→3`). Each function call gets a fresh memo. All parameter checks within that call share the same memo — that's how shapix knows `N=4` in `x` must match `N=4` in `y`.
 
-```python
-import shapix
-from beartype import beartype
+This happens **automatically** with `@beartype`. Shapix detects the beartype wrapper frame via `sys._getframe()` and associates a memo with it. No extra decorator needed.
 
-# Option 1: Memo only — pair with @beartype
-@shapix.check
-@beartype
-def f(x: F32[N, C]) -> F32[N, C]:
-    ...
+### `@shapix.check` (optional)
 
-# Option 2: Memo + beartype combined with custom config
-from beartype._conf.confmain import BeartypeConf
+`@shapix.check` provides **explicit** memo management. It's useful in two scenarios:
+
+1. **Combining memo with custom `BeartypeConf`** — a single decorator instead of stacking two:
+
+```python
+from beartype import BeartypeConf
 
 @shapix.check(conf=BeartypeConf())
-def f(x: F32[N, C]) -> F32[N, C]:
-    ...
+def f(x: F32[N, C]) -> F32[N, C]: ...
 ```
 
+2. **Exotic call stacks** where frame-based detection is unreliable (deep decorator chains, recursive wrappers).
+
+For normal usage, just use `@beartype` — it works out of the box.
+
 ### Manual checks with `check_context`
 
 For `isinstance`-style checks outside of decorated functions, use `check_context` with beartype's `is_bearable`:
@@ -383,7 +450,7 @@ Shapix uses three key mechanisms:
 
 1. **`Annotated[T, Is[validator]]`** — Each array type annotation (e.g., `F32[N, C]`) produces a `typing.Annotated` type with a beartype `Is[...]` validator. This lets beartype handle all the dispatch natively.
 
-2. **Frame-based memo management** — beartype's `Is[validator]` call stack is deterministic: `validator → _is_valid_bool → beartype_wrapper`. All parameter checks for one function call share the same wrapper frame. Shapix identifies this frame via `sys._getframe()` and associates a dimension memo (name → size bindings) with it. This is how cross-argument consistency works with zero boilerplate.
+2. **Frame-based memo** — beartype's `Is[validator]` call stack is deterministic: `validator → _is_valid_bool → beartype_wrapper`. All parameter checks for one function call share the same wrapper frame. Shapix identifies this frame via `sys._getframe()` and associates a dimension memo (name → size bindings) with it. This is how cross-argument consistency works with zero boilerplate.
 
 3. **Thread-local storage** — Each thread gets its own memo stack via `threading.local()`, ensuring thread safety.
 
@@ -453,6 +520,7 @@ def f(x: F32[~B, C]) -> F32[~B, C]:  # type: ignore[reportInvalidTypeForm]
 | BeartypeConf | Not supported (decorator conflict) | Fully supported |
 | Type checker | Metaclass magic (confuses pyright) | `Annotated` aliases (clean) |
 | Backends | NumPy, JAX | NumPy, JAX, PyTorch |
+| PyTree | Built-in with structure binding | Built-in with structure binding (via optree) |
 | Dependencies | jaxtyping + beartype | beartype only |
 | Custom decorator | Required | Not required |
 

pyproject.toml

@@ -5,15 +5,16 @@ description = "Add your description here"
 readme = "README.md"
 authors = [{ name = "acecchini", email = "ale.cecchini.valette@gmail.com" }]
 requires-python = ">=3.12"
-dependencies = ["numpy>=2.2.0", "beartype>=0.22.9"]
+dependencies = ["beartype>=0.22.9"]
 
 [dependency-groups]
-# CPU-only variants for dev/CI testing. Shapix is device-agnostic — it checks
-# .shape and .dtype only, so GPU/TPU builds (torch+cuda, jax[cuda12], etc.)
-# work without any extra configuration. Users install their own backend.
-torch = ["torch>=2.6.0"]
-jax = ["jax[cpu]>=0.5.0"]
+# Backends are CPU-only for dev/CI. Shapix is device-agnostic — it checks
+# .shape and .dtype only. Users install their own numpy/jax/torch/optree.
 dev = [
+  "numpy>=2.2.0",
+  "torch>=2.6.0",
+  "jax[cpu]>=0.5.0",
+  "optree>=0.14.0",
   "ruff>=0.15.4",
   "pyright>=1.1.408",
   "mypy>=1.19.1",

src/shapix/__init__.py

@@ -20,10 +20,14 @@ def conv(x: F32[N, C, H, W]) -> F32[N, C, H, W]: ...
 Exports
 -------
 Dimension symbols
-    ``B``, ``N``, ``P``, ``L``, ``C``, ``H``, ``W``, ``T`` — named dimensions.
+    ``B``, ``N``, ``P``, ``L``, ``C``, ``H``, ``W`` — named dimensions.
     ``__`` — anonymous (match any single dim, no binding).
     ``Scalar`` — scalar (no dimensions).
 
+PyTree structure symbols
+    ``T``, ``S`` — named tree structure symbols.
+    :class:`Structure` — create custom structure symbols.
+
 Unary operators (apply to any dimension)
     ``~N`` — variadic (match zero or more contiguous dims).
     ``+N`` — broadcastable (size 1 always matches).
@@ -32,6 +36,7 @@ def conv(x: F32[N, C, H, W]) -> F32[N, C, H, W]: ...
 Classes
     :class:`Dimension` — create custom dimension symbols with arithmetic support.
     :class:`DtypeSpec` — describe a set of allowed dtypes by string name.
+    :class:`PyTree` — subscriptable pytree annotation (requires ``optree``).
 
 Functions
     :func:`make_array_type` — create subscriptable array type factories for
@@ -46,6 +51,10 @@ def conv(x: F32[N, C, H, W]) -> F32[N, C, H, W]: ...
 from ._array_types import make_array_type as make_array_type
 from ._decorator import check as check
 from ._decorator import check_context as check_context
+from ._pytree import PyTree as PyTree
+from ._pytree import S as S
+from ._pytree import Structure as Structure
+from ._pytree import T as T
 from ._dimensions import __ as __
 from ._dimensions import B as B
 from ._dimensions import C as C
@@ -55,6 +64,5 @@ def conv(x: F32[N, C, H, W]) -> F32[N, C, H, W]: ...
 from ._dimensions import N as N
 from ._dimensions import P as P
 from ._dimensions import Scalar as Scalar
-from ._dimensions import T as T
 from ._dimensions import W as W
 from ._dtypes import DtypeSpec as DtypeSpec

src/shapix/_array_types.py

@@ -80,6 +80,7 @@ def __call__(self, obj: object) -> bool:
     # the memo with partial bindings from a bad argument).
     single_snap = memo.single.copy()
     variadic_snap = memo.variadic.copy()
+    structures_snap = memo.structures.copy()
 
     result = check_shape(tuple(shape), self._shape_spec, memo) == ""
 
@@ -88,6 +89,8 @@ def __call__(self, obj: object) -> bool:
       memo.single.update(single_snap)
       memo.variadic.clear()
       memo.variadic.update(variadic_snap)
+      memo.structures.clear()
+      memo.structures.update(structures_snap)
       self._fail_id = obj_id
 
     return result

src/shapix/_dimensions.py

@@ -53,7 +53,6 @@ def flatten(x: F32[N, C]) -> F32[N * C]: ...
   "C",
   "H",
   "W",
-  "T",
   # Anonymous
   "__",
   # Special
@@ -198,7 +197,6 @@ def _dim_spec(self) -> DimSpec | None:
   C: Dimension
   H: Dimension
   W: Dimension
-  T: Dimension
   __: Dimension
 else:
   # Common named dimensions
@@ -210,7 +208,6 @@ def _dim_spec(self) -> DimSpec | None:
   C = Dimension("C")
   H = Dimension("H")
   W = Dimension("W")
-  T = Dimension("T")
 
   # Anonymous (match anything, no binding)
   __ = Dimension("__")

src/shapix/_memo.py

@@ -32,6 +32,9 @@ class ShapeMemo:
   variadic: dict[str, tuple[bool, tuple[int, ...]]] = field(default_factory=dict)
   """Variadic dimension bindings: ``{"spatial": (False, (28, 28))}``."""
 
+  structures: dict[str, object] = field(default_factory=dict)
+  """PyTree structure bindings: ``{"T": <PyTreeSpec>}``."""
+
 
 _local = threading.local()