tenet.pytree

JAX pytree registration: importing this module makes SymmetricTensor a pytree.

tenet.pytree

Opt-in JAX PyTree registration. import tenet.pytree; core never imports this.

tenet.enable_jax is the one-call spelling of that import, and is what the docs teach; this module stays importable directly and its effect is unchanged.

The contract: the leaves of a SymmetricTensor are its data -- one dense matrix per coupled sector, in map_layout(structure).sectors order; the aux data (the treedef) is its structure, which is frozen, hashable and array-free (invariant 8) and therefore a sound JIT cache key.

That is a handful of large leaves rather than tens of thousands of small ones, which is a strict improvement for JAX: the per-block slices whose VJPs used to produce thousands of single-primitive XLA compilations are gone, and jax.tree.map over a tensor now walks the coupled sectors.

_unflatten deliberately does not validate, because JAX calls it with sentinels, tracers and whatever a jax.tree.map returned (https://docs.jax.dev/en/latest/custom_pytrees.html). The public constructor remains the trust boundary: SymmetricTensor(T.structure, T.blocks) re-runs the full check.

Batching recipe::

batched = jax.tree.map(lambda *bs: jnp.stack(bs), T1, T2, T3)
jax.vmap(lambda T: tenet.norm(T) ** 2)(batched)

batched is a transport container, not a tensor: its block shapes do not match its structure and nothing in tenet.* should be called on it. Inside the vmapped function the leaves are BatchTracers whose .shape is the unbatched shape, so ops and their validation behave normally.

Complex-dtype gradients are unverified: norm goes through abs, whose JAX derivative needs care for complex input. Tests here are real float64.