uniform_tucker_tensor_train#

Uniform Tucker tensor trains (UT3): a stacked-supercore + boolean-mask representation of (a stack of) ragged TuckerTensorTrain s, laid out for GPU/jit efficiency.

A UT3 pads the cores of a 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 records the real extent of each padded edge with boolean masks. It is, by design, a faster representation of the ragged layer – see docs/uniform_equivalence_contract.md and the other docs/uniform_*.md notes.

NOTE: this module is being built incrementally (slice 1: foundation). Linear algebra, orthogonalization, SVD, sampling, and the jax pytree registration land in later slices.

Classes#

UT3Masks

The static rank structure of a uniform Tucker tensor train: its two boolean edge masks.

UniformTuckerTensorTrain

A uniform Tucker tensor train: two supercores (the data) + a UT3Masks holder (the

UT3Weights

Diagonal weights on the internal edges of a UniformTuckerTensorTrain -- the uniform

Functions#

ut3_absorb_weights(x, weights)

Contract diagonal edge weights into x's supercores (shape-preserving): the returned

ut3_weighted_norm(x, weights[, use_orthogonalization])

Weighted Hilbert-Schmidt norm ||absorb_weights(x, weights)|| (shape stack_shape; a scalar

ut3_weighted_inner(x_A, weights_A, x_B, weights_B[, ...])

Weighted Hilbert-Schmidt inner product