# Why uniform T3s let ranks vary across the stack (but not shape) > A design-philosophy note for `UniformTuckerTensorTrain` and the `ut3_*` backend. It records *why* > we let the ranks of a stacked uniform T3 differ from one stack element to the next — and why the > physical shape is, by contrast, locked across the stack. The aim is the reasoning, not a rulebook: > if you understand the manifold-vs-variety picture below, the rules fall out on their own. --- ## The representation, and what the mask actually means A `UniformTuckerTensorTrain` pads a (possibly ragged) Tucker tensor train up to common sizes `n` (Tucker rank), `r` (TT rank), `N` (mode dimension), stacks the `d` cores onto a leading axis to form *supercores*, and carries boolean **edge masks** marking the real extent of each padded edge. The first thing to be clear about: **the mask does not define the dense tensor — the zeros do.** A stack element whose Tucker rank is `n_s < n` simply has zero columns in the padded slots, so it represents a tensor of rank `≤ n_s` regardless of the mask. Masking-down and "a rank-`n` core that happens to have zero columns" are the same tensor. What the mask *does* define is the **rank** — equivalently, *which stratum of the rank hierarchy the point sits on*. That is exactly the information that rank-aware operations need: where SVD truncation should cut, how `ut3_to_t3` extracts the real sub-block, whether the ranks are minimal. So the mask is per-element metadata: each stack element carries its own "official rank," static and known at trace time (which is why masks ride along as static aux data under `jit`). ## The real question: a batch on a manifold, or a batch in a variety? Because the mask is a stratum label, "should ranks vary across the stack?" is really: > Is a stacked uniform T3 a batch of points on **one smooth manifold**, or a batch of points in **one > bounded-rank variety**? Both have a clean mathematical home: - **Fixed rank (uniform across the stack).** The set of T3s with ranks *exactly* `(n, r)` is a smooth **manifold** `M₍ₙ,ᵣ₎`. This is the home of everything Riemannian — tangent spaces, retraction, gauge, the Taylor-series machinery the library is named for (T4S). A uniform-rank stack is a batch of points on a single such manifold. - **Varying rank (padded to `(n, r)`, one mask per element).** Tucker ranks are matricization ranks and TT ranks are unfolding ranks — both ordinary *matrix* ranks — so the set of tensors with ranks **≤ `(n, r)`** is a **determinantal variety** `V₍≤ₙ,≤ᵣ₎`, cut out by the vanishing of unfolding/ matricization minors. It is the closure `⋃₍ranks ≤ (n,r)₎ M₍ranks₎`, stratified by the lower-rank manifolds and singular along them. The padded-supercore-plus-mask representation is *exactly* a parametrization of a point in this variety, with the mask naming the stratum. We choose the **variety**. The decisive reason is a use case the manifold cannot serve well: **ideal ranks are rarely known a priori.** Sweeping a batch of candidate ranks at once — to see which performs best on a task — is a batch of points at *different* ranks, i.e. a batch in the variety. The mask layout expresses this for free (the padded supercore shapes stay fixed, so `jit` is unbothered), and it is a strict superset of the fixed-rank case: you can always restrict to one stratum when you want the manifold. ## Why shape is locked, and the principled asymmetry Shape is treated oppositely — it is **fixed across the stack** — and the contrast is not arbitrary: - **Shape determines the *ambient space*** `ℝ^{N₀×…×N_{d-1}}`. Stacking means "a batch of objects living in the *same* space." Let the shape vary and `to_dense` could not even return one `stack_shape + (N₀, …)` array; "stack" would lose its meaning. - **Rank determines *model complexity within* that ambient space** — which point of the variety, which stratum. That is free to vary, because every element still lives in the one shared space. So shape consistency is a **structural invariant** (a hard error if violated), while a varying rank is just a different point of the same variety. The asymmetry is the right one: same ambient space, different models inside it. (One nuance, since "shape fixed" is easy to misread: it means fixed *across the stack*. Shape still varies *across modes* — `N₀,…,N_{d-1}` differ — which the supercore handles by padding each mode to `N=max(Nᵢ)` plus `shape_mask`. That cross-*mode* padding is a separate axis from the cross-*stack* invariant; only the latter is the structural hard error.) ## Relationship to ragged `TuckerTensorTrain`, and the one real cost A ragged stacked `TuckerTensorTrain` *physically cannot* carry varying ranks — its cores are single-shape arrays, so uniform rank is forced on it. `UniformTuckerTensorTrain` exists precisely because the mask layout can express more. The relationship is clean and worth keeping in mind: - uniform-rank UT3 stacks ↔ ragged stacked T3 (a faithful round-trip), and - UT3 *additionally* represents varying-rank stacks that a single ragged stacked T3 cannot. This is exactly why the **structural** rank checker is shaped differently in the two layers: `has_minimal_ranks` is a **scalar in ragged** (one core shape ⇒ ranks shared across the stack, so there is nothing per-element to report) but **per-element (shape `stack_shape`) in uniform** (ranks vary across the stack). The *numerical* checkers (`is_orthogonal`, etc.) are per-element in both layers — see `docs/batching_and_stacking.md` §9. The honest cost shows up at exactly one place — the ragged round-trip: 1. **Round-trip.** `ut3_to_t3` of a varying-rank stack can only return a *tree* of distinct T3s, not one stacked `TuckerTensorTrain` — which is the truthful answer, since the elements really do have different ranks. (`to_dense` is unaffected: the ambient shape is uniform across the stack.) ## The tangent layer keeps the varying-rank stack — uniform rank is a `K`-stack property, not a `C`-stack one It is tempting to think a *tangent* forces uniform rank across the stack ("a tangent needs a single tangent space, so the stack must be one smooth manifold point"). That conflates a **single** tangent vector with a **stack** of them. The distinction is `K` (tangent stack) vs `C` (frame stack); see `docs/batching_and_stacking.md`: - A single tangent vector — or a **tangent (`K`) stack**, a bundle of tangents sharing *one* frame — lives in **one** tangent space, so its rank is uniform. This is automatic (they share a frame) and is the *one* genuine uniform-rank requirement. The bv-pair check enforces it structurally: the variation rank masks must be **constant along `K`** (= the frame's gauge-shifted masks broadcast over `K`). - A **frame (`C`) stack** is a batch of tangents at *different* bases. It is a tangent vector to the **product** manifold `M_{n_1,r_1} × … × M_{n_C,r_C}`, which is smooth even when the factors are different fixed-rank manifolds; the per-element tangent spaces may have different dimensions and nothing couples them. So **ranks may vary across `C` exactly as in the plain layer** — the rank-sweep use case (a batch of candidate models at different ranks, compared against each other) is a varying-`C` stack of tangents. The price of varying `C` is a **discipline, not a precondition**: every tangent / manifold op is per-element-mask-aware (it must be anyway for the masks to mean anything), so it vectorizes over `C` with each element applying its own mask — `tangent_space_dimension`, `inner`/`norm`, gauge, retraction are all per-element. The one subtle spot is a batched Riemannian solve (Gauss-Newton / CG) over a varying-rank `C` stack: the masked-out directions must sit in `ker(J)` so they contribute zero to the CG reductions (`rᵀr`, `pᵀAp`) and never pollute the active per-element solve — which holds as long as the gauge projection and `J`/`Jᵀ` respect the masks. Varying-`C` is a first-class case in the uniform-tangent equivalence-contract tests. ## Why there is no `to_vector`/`from_vector` on the uniform tangent `T3Tangent` offers a flat-vector interface (`to_vector`/`from_vector`, plus `save`/`load`); `UT3Tangent` deliberately does not, and this section is why (it is a consequence of the variety view above, not an unfinished port). A varying-rank stack has **no single flat-vector reconstruction signature**: `from_vector(vec, shape, ranks)` cannot encode per-element ranks in one signature. And both interop paths bypass the need: jax-native optimization works on the **pytree directly** (the padding has zero gradient, so nothing is gained by flattening), while flat-vector interop with e.g. scipy routes through the **ragged layer** (`ut3_to_t3` → `to_vector`), which is clean for the uniform-rank stacks scipy-style optimizers handle anyway.