Weighted layer — implementation internals#

Design records for the weighted layer (docs/weighting.md is the user-facing account). Two builds produced it: the ragged layer and then its uniform mirror, and most of what is worth preserving is why a tempting alternative is wrong, plus three corrections the builds made to their own plans.

Two classes, because a tensor and a tangent have different edges#

T3Weights pairs with a TuckerTensorTrain; T3FrameWeights pairs with a tangent. They are not the same object wearing two hats:

A T3 as a network has exactly two internal edge families — Tucker-rank edges nᵢ (d of them) and TT-bond edges rᵢ (d+1, ends trivial). So a weight is (tucker[d], tt[d+1]) — which is exactly the shape t3svd returns, so the singular values are literally a T3Weights.
A tangent is not one network. It is a sum of d diagrams, each with one variation core, and the edges left of the variation are left-orthogonal (rL) while those right are right-orthogonal (rR). Those four rank families (nU/nD/rL/rR) are not a labelling choice — they are which side of the variation an edge is on, and the shapes forbid mixing (rLₖ ≠ rRₖ in general; a left diagonal will not fit a right bond).

Forcing one class would need down=None, left=right for the plain case — uglier than two honest types.

The tangent weight is a METRIC (the design change that caught us out)#

The first design absorbed the tangent weights into the frame cores (up→U, down→O, left→P, right→Q, variations untouched). It shipped as the opposite: a metric on the variation coordinates, absorbed into V/H, with the frame left orthonormal.

Why: the frame’s d+1-th left/right cores are base-point padding — “not really part of the frame” (docs/frame_variations.md) — so there are only d natural tangent edges, matching the d variation cores. And the metric reading is what Grasedyck–Kramer actually needs.

The abandoned design is worth knowing about because it explains a piece of machinery that does not exist: absorbing into the frame breaks orthonormality, which would force the weighted norm through the doubled-rank tv_to_t3 embedding. Under the metric design the frame stays orthonormal, so the weighted norm is just the corewise norm of the weight-scaled variations — O(ranks), no doubled tensor, no tv_weighted_* twins. If you find a note describing that path, it is pre-change.

A weight batches with the FRAME but is absorbed into the VARIATIONS#

This is the subtlety, and it is the one that actually bit. The two questions are independent:

question	answer for `T3FrameWeights`
what does it act on?	the variations (`down`→V, `up`/`left`/`right`→H)
what does it batch with?	the frame — stack `C`, never the variations’ `K + C`

A metric is one-per-base-point, shared by all K tangent vectors at that frame. At the design change, the implementation moved the absorption target to the variations — correctly — and moved the stack along with it, which was wrong. Nothing computed the wrong number (the leading '...' broadcasts C over K + C for free, and a C-only weight was verified bit-identical to a K-tiled one), but fv_weights_consistent demanded the weight match the variations’ full K + C stack, and so rejected the canonical Grasedyck–Kramer weight — from_t3weights(from_t3svd(x)) is C-stacked — the moment there was a K-stack of tangents. The predicate contradicted the operation it described.

It survived because the tests only ever built K+C weights, so the broken case was never constructed, and because is_consistent_with is purely diagnostic (in no enforcement path) — a lying predicate, not a wrong result. Fixed 2026-07-15; the lesson is in the table above.

The two-level check that came out of it#

Pairing a weight with variations and pairing it with a frame are different questions, so they get different checks:

fv_weights_consistent(variations, weights) — the trailing rule, blind to the frame. The weight’s stack must be the trailing part of the variation stack. This is exactly the rule check_fv_pair already applies to (frame, variations), with the weight playing the frame’s role. It stays blind on purpose: the variations are intentionally blind to their frame, and the standalone fv_absorb_weights(variations, W) has no frame to consult.
check_fw_pair(frame, weights) — the exact stack, where the frame is present. A K+C weight passes the trailing rule too (it reads as C_w = K+C: that many base points, one tangent each — a coherent object, just not this frame’s metric), so only here is there enough information to reject it. Wired into T3Tangent.weighted_norm/weighted_inner/absorb_weights.

Rejected: checking the weight against the frame in the backend predicate (fv_weights_consistent(frame, weights)). It invents a second pairing pattern where the library already has one, and couples the predicate to a frame the standalone absorb does not have.

The uniform mirror#

A uniform weight is “the ragged weight, padded for performance.” It carries the edge masks of the object it weights (a weight’s edges are the tensor’s edges) and has no shape field — weights live only on internal edges; external/physical-mode weights are out of scope.

`absorb` is garbage-transparent — it needs no masking#

The pre-build plan asserted that “absorb and norm/inner are reductions, so garbage padding must be zeroed first.” True of norm/inner; false of absorb, and the claim is load-bearing enough that a reader would implement a pointless entry-mask against it.

absorb is a pointwise scale along each edge axis, not a reduction: output slot i depends only on the core’s slot i and the weight’s slot i, so garbage can never mix into a real slot. Garbage in the padding gives garbage in the padding, which the equivalence contract declares don’t-care. The reductions that follow (weighted_norm/_inner) reduce to the existing plain uniform norm/inner, which mask their own inputs. Same for concatenate/kronecker: supercores are copied or outer-product’d, never reduced.

The masking/weighting wall, and why it cost nothing#

Masking and weighting are kept apart in code even where the mechanics coincide, because they are opposite kinds of thing: masks are boolean structure (static aux, non-differentiable, define rank, value-neutral); weights are float parameters (traced leaves, differentiable, scale the contraction, value-affecting). This is the same distinction that uniform_rank_masks_rationale.md invokes to reject “a float mask doubling as edge weights.”

Weighting code never calls a *_apply_masks / *_make_masks function. Where the mechanics genuinely coincide, the shared core is broken out into a neutral subfunction that each side calls.

Because absorb needs no masking, the wall falls out rather than being imposed — backend/ut3_operations.py does not import the masking layer at all. The rule forced exactly two extractions, both of which improved the layer independently of weighting:

common.prefix_mask(ranks, pad) — arange(pad) < ranks[..., None] was written out in ~8 places across four modules, including a private _prefix_mask the mask builders themselves did not use.
common.require_concrete_masks — already documented as “a jit guard (infrastructure), unprefixed by design”, yet living in ut3_masking; ufv_masking never used it, so the home was accidental. It belongs beside ValueHashedMasks and prefix_mask: those three are the uniform mask-representation contract (concrete host arrays · value-based hash/eq · canonical prefix form), obeyed by every uniform object’s masks rather than owned by any one family.

(That move also surfaced a real gap, since fixed: the whole frame/variation side never called the guard — ufv_apply_frame_masks, ufv_apply_variations_masks, ut3frame_to_t3frame and ut3variations_to_t3variations all accepted traced masks and failed with jax’s cryptic TracerArrayConversionError from inside a numpy call, rather than the message that names the fix. All four now guard, and tests/test_dispatch.py::TestTracedMaskGuard pins every chokepoint across all four uniform object kinds — the plain-layer guard previously had no test at all, only a doctest.)

`reciprocal` must guard the padding — and this is why weights carry masks#

A canonical weight’s padding is zero, and 1/0 = inf. Masking downstream works by multiplication, so 0 × inf = nan then poisons every masked reduction — the exact hazard ../uniform_equivalence_contract.md names (“the fill must be finite”). This is not a corner case: the Grasedyck–Kramer metric is from_ut3svd(x).reciprocal(), so the headline path runs straight through it.

The fix is the standard double-where: neutralize the padding to 1, apply fn, then overwrite the padding with 0. Both halves earn their keep — the inner one keeps the gradient finite (a nan from a dead branch still propagates backward through a single where), and it neutralizes large-finite garbage padding too, so no separate entry-mask is needed. sqrt uses the same helper (sqrt is fine at 0 but not differentiable there, and garbage padding can be negative).

Real-slot zeros are deliberately NOT guarded. A genuinely zero singular value gives inf, exactly as in ragged. It is real data, not a padding artifact; silently clamping it would hide a rank-deficient point.

This is also the argument that a uniform weight must carry masks: a maskless weight cannot implement reciprocal correctly.

`concatenate` / `kronecker`: the mask is just another weight#

The pre-build plan billed these as “the hard ones” and “the most error-prone code in the build.” Wrong — and worth correcting, because it sends a reader looking for a difficulty that is not there.

Treat the mask as another weight that happens to hold 0s and 1s. What is needed is

weight_AB * mask_AB  ==  combine(weight_A * mask_A, weight_B * mask_B)

and it is satisfied by combining the weights and the masks the same way, because both concatenation and the Kronecker product commute with elementwise multiply (kron(a∘p, b∘q) = kron(a,b) ∘ kron(p,q) — the mixed-product property; nothing special to booleans). Each op is one operation applied twice, with no mask cleverness at all. That closure is the argument for boolean masks over integer ranks: an integer rank cannot even express the strided result.

The genuine risks are ordering/axis mistakes:

kron must be a last-axis outer product broadcasting the shared (d,)+stack prefix — NOT np.kron, which would Kronecker the mode/stack axes too. The ragged build hit exactly this.
A-major must agree with any core-combine that pairs with it. In uniform there is none yet — there is no ut3_mult, only ut3_add — so ut3_kronecker_weights ships verified against the ragged oracle but unpartnered. If ut3_mult is ever built, its index pairing must match.

The output masks are gappy (concat leaves a hole wherever an input had rank slack; kron’s real set {a*nB + b : mask_A[a] and mask_B[b]} is strided over the padded width, so even two prefix inputs give holes). That is a description of the result, not a difficulty: it costs the combines nothing, and only obliges consumers to read the mask instead of slicing a prefix. It cannot be flattened — slot a*nB + b with b >= rank_B holds wA[a] * <padding>, so a prefix mask of rank rA·rB would claim padding as real data (phantom rank).

Consuming a gappy mask needs no argwhere — boolean indexing (w[m]) selects the real slots in ascending order directly, and works on numpy, on jax eagerly, and even under jit (the masks are closed-over host constants, so jax resolves the shape at trace time; NonConcreteBooleanIndexError fires only for a traced mask, which require_concrete_masks forbids). The point was never “gather vs slice” — it is that a rank mask may be gappy, so the real slots must be selected through the mask rather than as a prefix. (The physical shape genuinely is a contiguous prefix, and does slice.)

The one precondition uniform adds#

Ragged catches a weight/object rank mismatch as a loud einsum shape error (a length-n weight vector against a rank-n core). Uniform pads both to a common (n, r), so a mismatched mask is invisible to the shapes and would silently corrupt: a weight whose mask calls slot i padding carries a canonical zero there, so absorbing it zeroes a real slot. Hence an explicit structural predicate (ut3_weights_consistent / ufv_weights_consistent) with frontend enforcement — the same precondition uniform adds to variation add/sub (../uniform_masks_vs_ranks.md).

Placement notes (both forced, not preference)#

utv_weighted_norm / utv_weighted_inner live in utv_operations, not beside ufv_absorb_weights as their ragged twins do (fv_weighted_norm is in fv_operations). They need utv_corewise_inner, and utv_operations already imports ufv_operations — the ragged placement would be a circular import. Ragged has no such constraint: its fv_ layer reaches a standalone corewise module.
The frontend free functions carry the family prefix (t3_absorb_weights / ut3_absorb_weights / fv_absorb_weights / ufv_absorb_weights). Ragged/uniform × tensor/tangent gives four distinct “absorb the weights” operations; unprefixed they would all be absorb_weights and could not coexist at the package root. See ../naming_conventions.md — the rule is that a free function has no class namespace, so the prefix does that job.

No ragged/uniform dispatch — the module IS the dispatch#

The weighted free functions are deliberately parallel and module-scoped: nothing infers the layer from the argument. The user controls dispatch by deciding which layer they work in, and the conversion hooks are how they switch.

optimizers.py’s isinstance(x0, UniformTuckerTensorTrain) is not a counterexample: newton_cg is a single entry point a user calls with whatever x0 they have, so it has no module to dispatch through. absorb_weights does.

What’s deferred#

Weighted + / − / scale / ⊙ as operations, and an optional thin container for operator sugar (a + b, a.norm()). Guaranteed reachable: every weighted op reduces to a move the substrate already provides — add → direct-sum cores + concatenate weights (+ tail-squash); ⊙ → Kronecker-combine cores + kronecker weights; inner/norm/apply/entries/probe → absorb then the plain op. So each is a pure addition (a new free function), never a refactor.
The Grasedyck–Kramer SingularValueRegularizer — the intended consumer. If M = W² then ⟨x, Mx⟩ = ‖absorb(x, W)‖² and Mx = absorb(x, W²), so absorb-into-tangent is exactly its primitive; it builds W from the frame’s singular values and calls absorb. Both layers now have everything it needs. See fitting_internals.md.
External / physical-mode weights (on the Nᵢ output legs) — out of scope, deliberately. The refactor made clear the external edges are of a different character than the internal ones; the old code’s shape_weights family was dropped.
Backend ut3_norm / ut3_inner twins — see deferred_and_rejected.md.

Build records#

dev/archive/weighted_layer_design.md (ragged) and dev/archive/uniform_weighting_design.md (the uniform mirror; §8 is its decision log). Both are historical — where they disagree with this note or with docs/weighting.md, they are wrong.