# Rank masks — the decision record (why they exist, and why boolean)

> The consolidated mask decision record, in two parts: **why the masks exist at all** (they cannot
> be dropped — they are the functional rank constraint that makes variable-rank optimization work),
> and **why the representation is boolean masks rather than integer ranks** (the closure argument,
> the honest ledger, and the rejected alternatives). The user-facing "how masks behave" story is
> `docs/uniform_masks_vs_ranks.md`. Pairs with `docs/uniform_ranks_and_varieties.md` (what the
> object represents), `numerical_contract_catalog.md` (the minimal-rank audit),
> `uniform_svd_prefix_orthogonalization.md`, and `uniform_pytree_composition.md`.

## Part I — why the masks exist (and why "just inflate to uniform rank" fails)

## What the masks do

A uniform Tucker tensor train pads each edge to a common size `R` — so a whole *stack* of T3s can live in
one array, `lax.scan`-friendly — and marks slots `[0, rank)` real with a prefix mask. The real rank may
differ **per edge and per stack element**: that is the *determinantal variety* — fixed shape, ranks that
vary (`docs/uniform_ranks_and_varieties.md`). In the tangent / optimization layer the masks **zero the
variation in the padded slots**, so a gradient step cannot put content beyond the real rank.

## Why variable ranks matter

Real data rarely wants one rank everywhere. A good T3 fit may need a small rank on one edge and a large
one on another; forcing a single rank either **underfits** (rank too low — accuracy left on the table) or
**overfits** (rank too high — fitting noise, wasting compute). *Rank continuation* exploits exactly this:
start low and grow ranks edge-by-edge under a conditioning criterion
(`docs/rank_continuation.md`, `examples/fit_varied_rank_tensor_newton_cg.py`). All of it presumes you can
hold ranks at chosen, possibly-varied values — i.e. that the optimizer **respects the rank structure**.

## The tempting simplification — and why it fails

The minimal-rank audit (`numerical_contract_catalog.md`, empirically verified against the dense oracle)
found that **no manifold operation requires minimal rank**: `inner`/`norm` need only an orthonormal,
gauged frame; `project`/`transport`/gauge need only orthonormality; `retract` is valid on a non-minimal
frame. Equivalently, for a **given** tangent vector, inflating its ranks to the padded size (orthonormally
completing the frame, zero-extending the variations) leaves every operation's result unchanged:

```
op(inflate(tangent)).to_dense()  ==  op(tangent).to_dense()
```

So why not drop the masks entirely, work on the padded supercores, and treat every slot as real? It would
be a large simplification — no static structure in the jit cache key, so the per-step-recompile problem
would vanish for the tangent layer.

**Because optimization does not operate on a fixed tangent — it computes a new one (the gradient) every
step.** At a rank-deficient frame the projected gradient generally **has content in the padded
("completion") directions** — moving there *increases* the rank. The masks are exactly what zero that
content. Drop them and the gradient grows every edge toward the common padded rank `R`: you lose rank
control and overfit — the precise failure the variable-rank feature exists to prevent.

So the maskless equivalence is genuine, but only for **operations on a fixed tangent**. The masks are not
bookkeeping; they are a **functional rank constraint** that makes variable-rank fitting work. Keeping them
is the deliberate choice.

## The cost of keeping masks, and how it is paid

Masks are static structure, so they ride in jax `aux_data` and become part of the `jit` cache key — which
*risks a recompile* every time the frame is rebuilt (i.e. every optimization step). Two design choices
remove that cost:

- **SVD-based orthogonalization** places the real content in the leading slots, so the masks come out a
  **deterministic prefix** — bit-identical every step at fixed rank
  (`uniform_svd_prefix_orthogonalization.md`).
- **Value-hashed mask holders** key the `jit` cache on mask *content*, so a rebuilt-but-identical holder
  is a cache hit (`uniform_pytree_composition.md`).

Together: **no recompile within a fixed-rank stage; a recompile only when the ranks genuinely change**
(e.g. a rank-continuation step) — which is correct, and rare.

## Part II — why boolean masks rather than integer ranks (the representation decision)

*(Excised from the pre-split `uniform_masks_vs_ranks.md`, whose user-facing successor carries the
"how masks behave" story; this part preserves the decision ledger.)*

### The two candidate representations

The rank metadata of a uniform T3 must be *stored*, not derived: a padded core cannot tell you its
rank (a "real" slot may be numerically zero; a padded slot may carry garbage). The question is the
*format*:

- **(a) Boolean edge masks** (chosen): for each edge, a 0/1 vector over the padded size marking
  which slots are real. Equivalently, a diagonal **projector** onto the real subspace.
- **(b) Integer rank arrays**: store the counts (`tucker_ranks`, `tt_ranks`) and rebuild a prefix
  mask `arange(pad) < k` on demand. This presumes the real block is always the **prefix** `[0, k)`.

These coincide *only* if every mask is prefix-contiguous — and the natural operations do not
preserve prefix-contiguity (the concat/Kronecker gap examples are in the user doc). The **boolean
algebra of subsets** (under disjoint union and product) is *closed* under exactly the operations
T3 algebra is made of; the **integers-as-prefixes** representation is not: every `+` and `×`
breaks the prefix invariant, and the only way to restore it is to gather/scatter the data to the
top-left. In the stacked setting that compaction is *ragged per element* — precisely the
irregularity the uniform layout exists to avoid. Read this way, the mask is "more minimal" in the
sense that matters: it is the **minimal state the operations preserve**. The integer rank is
minimal *per object*, but only stays valid if you continuously renormalize positions — i.e. pay
for data movement, or secretly re-derive the positions (reinventing the mask).

### Honest pros and cons

**Boolean masks (a):**
- ✅ Closed under ⊕/⊗ (concatenate / Kronecker); zero data movement. The composition niceties
  (add → concat masks, multiply → Kronecker masks, sweep → apply masks) are consequences, not
  coincidences.
- ✅ Directly usable in the masking einsums.
- ✅ Minimal in the operationally relevant sense (state the algebra preserves).
- ➖ The rank is implicit (`mask.sum(-1)`), so it reads less directly than an integer field.
- ➖ A mask can sit in a non-canonical or meaningless pattern that `validate` cannot cheaply vet —
  the same non-enforced status as orthogonal/minimal/gauged.
- ➖ Masks do **not** save you from rank growth (stated user-side too).

**Integer ranks (b):**
- ✅ Smallest per object; reads literally as the *stratum label* of the determinantal variety.
- ✅ Trivially the canonical form.
- ➖ **Not closed**: every `+`/`×` requires a per-stack-element compaction to restore the prefix —
  ragged data movement — or it must carry positions anyway, which is the mask.

### Other representations considered (and why they're off the map)

Each fails one of the required properties — **closure under ⊕/⊗**, **uniform shape across the
stack**, and (for jit) **host-static** (`uniform_pytree_composition.md`):

| representation | closed under +/× | uniform shape across stack | notes |
|---|---|---|---|
| **bool mask (dense array)** ✅ chosen | ✅ concat / Kronecker | ✅ fixed padded size | the working representation; rank reported as `mask.sum(-1)` |
| integer rank **counts** | ❌ gappy ⇒ per-element compaction | ✅ | a count can't denote a gap; compaction is ragged data movement *and* data-dependent shapes ⇒ not jittable |
| hot-**position** lists / "hot rank" tuples | ✅ position arithmetic | ❌ varying length ⇒ ragged | a *sparse* mask; per-element lengths differ ⇒ jagged, defeating the uniform layout (you'd pad+remask — reinvent the mask) |
| **bool tuples** (vs arrays) | ✅ | ✅ | value-hashable, but non-contiguous/boxed ⇒ must materialize for the multiply; slower eager, worse stacked memory; no win over a bool array (which can be value-hashed via its bytes) |
| **float** mask doubling as edge **weights** | — | — | conflates *structure* (static aux, non-differentiable, defines rank, value-neutral) with a *parameter* (traced leaf, differentiable, scales the contraction, value-affecting): opposite jax treatment, and autodiff would silently differentiate — and a grad step corrupt — the "mask." Weights are a separate (parked) concept |

The bool mask is the unique fixed point: **closed ∧ uniform-shaped ∧ host-static.** The one
attractive property of the integer/tuple forms — value-based hashing for jit-cache hits — is real
but *separable*: get it on the bool mask via a byte hash (`uniform_pytree_composition.md`).
(Caveat: a rank *count* is an insufficient cache key off canonical form — two gappy masks with
equal counts but different positions are different computations — so the value hash keys on the
mask bytes; in canonical form the count suffices.)

### Maintainer guard: host numpy is intentional

Whichever encoding, the structure must live as **numpy (host)** static `aux_data`: any jax op on
it inside a trace becomes a tracer. Under jit the host masks fold into the compiled program as
device constants (no per-call transfer); the eager cost and the deferred `jax.device_put` option
are in `uniform_pytree_composition.md`. **The `np.*` in mask code is intentional — do not "fix" it
to `xnp`.**