# Composing UniformTuckerTensorTrain: dynamic supercores + a static structure holder

> A design-philosophy note for `UniformTuckerTensorTrain`. It records *why* the class is composed as
> two supercore fields (dynamic, differentiable) plus one small *structure holder* field (static,
> non-differentiable), rather than as a flat record — and why that holder is **value-hashed by its mask
> content** (so a rebuilt-but-identical holder is the same `jit` cache key). As with the other notes,
> this is an honest accounting of the choice and its costs, not a
> rule to follow on faith. It pairs with `docs/uniform_masks_vs_ranks.md` (why the structure is stored
> as boolean masks) and `docs/uniform_ranks_and_varieties.md` (what the object represents).

---

## The forcing constraint: masks are static, and static lives in pytree `aux_data`

The supercores are data we differentiate, `jit`, and `vmap` through. The masks are *structure* — the
variety stratum plus ambient dimensions (`docs/uniform_ranks_and_varieties.md`) — and we explicitly do
**not** want gradients flowing through them. In a JAX pytree that distinction is exactly children vs.
`aux_data`:

- **children** are the dynamic leaves: traced by abstract shape/dtype, never hashed by content.
- **`aux_data`** is the static part, carried in the *treedef* — and the treedef is `jit`'s **cache
  key**, so `aux_data` must be **hashable**, and treedef equality compares it with `==`.

So masks belong in `aux_data`. But a raw array — numpy *or* jax — cannot be `aux_data`: it is
unhashable, and `array == array` returns an *array*, not a single bool. (An early UT3 draft put a
bare tuple of three mask arrays in `aux_data`; it did not survive `jit`.) Something has to
make the static masks hashable.

## The decision: a small structure holder, value-hashed by mask content

```python
@dataclass(frozen=True, eq=False)   # eq=False suppresses the dataclass __eq__ (it would do the
class UT3Masks(ValueHashedMasks):   # array-ambiguous `array == array`); the mixin supplies VALUE-based
    tucker_edge_mask                # __hash__/__eq__ over mask content (common.ValueHashedMasks)
    tt_edge_mask
    # exposes the rank metadata (tucker_ranks, tt_ranks, …) derived from the masks

@dataclass(frozen=True)
class UniformTuckerTensorTrain:
    tucker_supercore        # dynamic  -> pytree child
    tt_supercore            # dynamic  -> pytree child
    shape: Tuple[int, ...]  # static   -> pytree aux_data (a value-hashable int tuple — the physical mode
                            #             dims are a contiguous prefix, so an int tuple suffices, no mask)
    masks: UT3Masks         # static   -> pytree aux_data (value-hashed by content)
```

The holder hashes and compares by **mask content** — `common.ValueHashedMasks`: a cached `tobytes` hash +
`np.array_equal` equality, with an `is`-identity fast path. `eq=False` is still required (it suppresses
the dataclass-generated `__eq__`, which would compute the array-ambiguous `array == array`), so the
mixin's value-based methods stand.

**Why value, not identity.** The holder rides in `aux_data`, which *is* `jit`'s cache key. If it hashed by
*object identity* (the bare `eq=False` default), a rebuilt-but-array-identical holder would be a **new**
key → recompile. That is fatal for the layer's whole purpose: in a manifold-optimization loop the
orthogonal frame is rebuilt every iteration (`ut3_orthogonal_representations` produces fresh mask holders,
but at fixed rank the *structure is unchanged*), so identity hashing would recompile **every step**, and
recompilation dwarfs the per-step compute — the exact opposite of why the uniform layer exists. Value
hashing makes the cache key reflect the rank **structure**: identical structure → cache hit (compile
once); a genuinely different structure → recompile (correct). A regression test pins this contract
(`tests/test_dispatch.py::test_mask_rebuild_does_not_recompile`).

The holder is deliberately **not** itself a registered pytree: as opaque `aux_data` its masks are never
flattened into children, so they cannot be traced or differentiated — precisely the guarantee we want. It
is also genuinely meaningful, not just a hashability wrapper: it *is* the rank-structure descriptor
(`tucker_ranks`, `tt_ranks`, … derive from it), so the dataclass stays minimal (two rank masks; `shape`
and the sizes are not stored redundantly). The same mixin is reused by the frame/variation holders
(`UT3FrameMasks` and `UT3VariationsMasks`, in `uniform_frame_variations_format.py`).

## Consequences worth knowing

- **Masks are numpy (host), not jax — and this is load-bearing, not a free dtype choice.** The masks must
  be concrete host arrays because their *values* are read (`int(mask.sum())`), recomputed into aux, and now
  **value-hashed** (`tobytes`) as the jit cache key — all of which a tracer breaks. See *Masks are numpy
  (host) — the jit story* below.
- **The backend is unaffected (raw arrays).** The holder is a frontend/pytree concern; backend functions
  take **raw arrays** in a layout that mirrors the fields — `.data = (tucker_supercore, tt_supercore,
  shape, (tucker_edge_mask, tt_edge_mask))`, supercores flat, then the static `shape` int tuple, then the
  two rank masks grouped as a sub-tuple (supercore-only ops take `.data[:2]`; `shape` is `.data[2]`;
  mask-using ops unpack `.data[3]`). A user on raw `.data` is never forced through the holder — consistent
  with the backend/frontend razor.

## Masks are numpy (host) — the jit story

The masks are `aux_data`, i.e. *static structure*. The hard-won point (it cost a debugging session —
see `dev/archive/uniform_slice_handoff.md`) is that **the masks must be stored *and computed* as numpy (host)
arrays, even when the supercores are jax.** This is required for jit correctness; it is not a
backend-agnosticism slip.

**Why jax masks break under jit.** Inside a jit trace, *every* `jnp` op returns a tracer — even on a
concrete constant. So if the masks are jax arrays, any mask op inside a traced function is a tracer:

- `int(mask.sum())` — the host-int shape/rank extraction (e.g. `ut3_to_dense`'s static prefix-slice,
  `t3svd`, the `.shape` property) → `ConcretizationTypeError`: you cannot pull a Python `int` from a
  tracer.
- mask *recomputation* — the rank recurrences in orthogonalization/svd and the `+`/`×` concat/Kronecker
  → the new masks are tracers, which would then **leak into the output's `aux_data`**. (Value hashing now
  makes this fail fast — `tobytes` on a tracer raises at hash time — instead of silently producing an
  object whose masks are escaped tracers; `require_concrete_masks` guards it explicitly regardless.)

Numpy masks avoid both: numpy ops run on the host and are never staged into the jaxpr, so
`int(mask.sum())` is a real host int and recomputed masks stay concrete. **So all mask logic uses `np`,
not `xnp`:** the structure is resolved on the host; only the data (supercores) flows through `xnp`.

**Under jit this costs nothing — it is optimal.** Because the mask logic is numpy on concrete inputs,
the whole structural computation (rank recurrences, prefix masks, shape extraction, concat/Kronecker)
runs **once at trace/compile time on the host** and folds into the compiled program as constants. Per
call, only the *data* computation runs on the device, with the masks already resident as compile-time
device constants. That is the ideal split — structure at compile (host), data at runtime (device) —
with **zero per-call host↔device mask transfer**.

**Eager (non-jit) GPU is the only place a transfer happens, and it is small.** Run eagerly, the masking
multiply (`supercore * mask`) moves the (boolean, KB-scale) masks host→device per op — a latency-bound
~µs cost. It is not a new bottleneck: eager jax is already dominated by per-op Python *dispatch*
overhead (tens of µs/op — the reason to jit), which dwarfs the mask transfer; and terminal ops
(`to_dense`, scalar `inner`/`norm`, measurement results) pull a larger result back to host anyway. The
only way to avoid even this is device masks — which break jit. That is a bad trade: sacrificing the
entire jit path (uniform's whole purpose) to save µs on the off-purpose eager path (eager work is the
ragged layer's job).

**Deferred future option (do NOT build speculatively):** if eager-GPU uniform is ever *profiled* as a
real hot path, cache a device copy of the masks (`jax.device_put` once, reuse the handle) so repeated
eager ops don't re-transfer. Local optimization, not a design change, and explicitly deferred — the
standing guidance is "if you care about performance, jit," at which point the masks are free.

### How to jit a uniform op — purely functional, or via the frontend

The rule is just *keep the masks static, trace the supercores* — and the **purely functional path needs
no class** (the OO-averse backend user is fully supported):

- **Functional (raw `.data`).** Close over the masks (host constants) and trace only the supercores —
  the standard jax idiom:

  ```python
  masks = (shape_mask, tucker_edge_mask, tt_edge_mask)        # HOST bool, static
  dense = jax.jit(lambda tk, tt: ut3_to_dense((tk, tt, masks)))(tucker_sc, tt_sc)
  ```

  The closed-over numpy masks become compile-time device constants; the supercores are the traced args.
  (`static_argnums`/`static_argnames` is *not* the route — the masks are arrays, hence unhashable, so
  they can't be marked static without a wrapper; close-over is the clean move.)
- **Frontend.** If you do use `UniformTuckerTensorTrain`, `jax.jit` over it works with nothing to think
  about: the masks ride as `aux_data` (static), the supercores are children (traced).

**The misuse, and the guard.** Passing the masks *among the traced args* (e.g. `jax.jit(ut3_to_dense)(data)`,
which traces every leaf of the tuple, masks included) makes them tracers → the host-int extraction and
mask recomputation fail. The backend **guards** this: a traced mask raises a clear, actionable error
("uniform masks must be concrete host arrays … close over the masks and trace only the supercores")
instead of jax's cryptic `ConcretizationTypeError`. So the failure is self-explaining; the right and
wrong forms are shown as a doctest.

### Why `T3Frame` (jax arrays) is fine as aux, but masks aren't

`T3Tangent`'s `T3Frame` may hold jax arrays as aux (its arrays are used only as data; in the current code
it actually flows as a pytree *leaf* via a numerical same-frame guard — see `manifold.py`). The UT3 masks
cannot be jax — for reasons that are about *what the op does with the aux*, not the dtype label:

- **`T3Frame` is aux used only as *data*.** Its cores feed contractions (`einsum`), whose tracer results
  flow on as data. It is never read as a host value (its *structure* is its array shape, which jax tracks
  statically — no `int(frame…)`), and a tangent op at a fixed frame **reuses the same input frame**
  unchanged as the output aux. Nothing forces it concrete → jax is fine.
- **The masks' *values* ARE host-readable static structure — read, recomputed, AND now hashed.** A uniform
  op (a) pulls the real shape/ranks as host Python ints — `int(mask.sum())` — to slice the padded supercore
  (the mask *values*, not the array shape, say which slot is real); (b) **recomputes** the masks (ranks
  change under `+`/`×`/svd) into the output's *static* aux; and (c) the holder now **value-hashes** the
  masks (`tobytes`) as the `jit` cache key. All three demand concreteness → host numpy (a jax mask gives a
  tracer on `.sum()`/`.tobytes()`, and recomputed jax masks would leak into aux). Value hashing actually
  *strengthens* this: a stray tracer mask now fails fast at hash time rather than leaking silently.

One line: **frame = aux-used-as-data** (jax fine); **masks = aux whose values are host-readable structure
and are recomputed** (must be numpy). (The frame *would* hit the same wall if you re-orthogonalized a frame
inside jit and wrapped a fresh tangent around it — which is why the convention is to hold the frame object
stable; uniform can't sidestep it, because changing the structure *is* the op.)

> **⚠️ Maintainers (human or AI): the `np.*` in the uniform mask code is INTENTIONAL.** Historically a
> bare `np.` was a tell that code wasn't backend-agnostic (should be `xnp`). **That heuristic does not
> apply to mask logic.** Masks are host structure and MUST be numpy. Do **not** "fix" mask `np.*` to
> `xnp`/jax for consistency — it silently breaks jit (tracer-leak into `aux_data`, `int()`
> concretization errors). The rule across the uniform layer is: **supercores (data) → `xnp`; masks
> (structure) → `np`.** (`ut3_make_masks` and the other mask builders therefore always emit numpy,
> with no `use_jax` flag; `to_jax`/`to_numpy` convert the supercores only, never the masks.)

## Honest costs

- **One extra class and a composition step**, versus a flat five-field record. Mild, and offset by the
  holder being a real object, but real.
- **Value hashing costs a small content hash/eq per `jit` call** (the price of *not* recompiling). The
  holder hashes its mask bytes (`tobytes`, cached on the frozen object) and compares with
  `np.array_equal` (with an `is`-identity fast path). That is `O(mask size)` = rank-sized — far smaller
  than the supercores, and negligible next to the op, let alone the recompile it prevents. We key on the
  **mask bytes**, not an integer rank *count*: off canonical form the real slots scatter (add = concat of
  two prefix masks → a gappy, non-prefix pattern), so a rank count cannot key it (see
  `uniform_masks_vs_ranks.md`); the bytes key both canonical and gappy forms. *(Optional further
  optimization: have the optimizer reuse the same `masks` object across iterations at fixed rank, so even
  the `array_equal` short-circuits on `is` — gravy, not required for correctness.)* This **retires** the
  earlier identity-hashing caveat (the original design deferred value hashing; the optimization-loop
  recompile it caused is exactly why we no longer defer it).

## Scope note

The pytree registration and value hashing described here are implemented and pinned by regression
test (`tests/test_dispatch.py::test_mask_rebuild_does_not_recompile`). This note remains the record
of *why* the class shape is `supercores + structure holder` — the decision that set the constructor
signature, the `.data` layout, and every derived property.