ut3_make_masks#

t3toolbox.backend.ut3_masking.ut3_make_masks(tucker_ranks, tt_ranks, n, r)#
def ut3_make_masks(
        tucker_ranks:   NDArray,            # HOST int, shape=(d,)   + stack_shape
        tt_ranks:       NDArray,            # HOST int, shape=(d+1,) + stack_shape
        n:              int,                # padded Tucker rank,     n >= max(tucker_ranks)
        r:              int,                # padded TT rank,         r >= max(tt_ranks)
) -> typ.Tuple[
    NDArray,  # tucker_edge_mask, HOST bool, static, shape=(d,)   + stack_shape + (n,)
    NDArray,  # tt_edge_mask,     HOST bool, static, shape=(d+1,) + stack_shape + (r,)
]:

Build the prefix RANK edge masks for a uniform Tucker tensor train.

Slot j of an edge is marked real iff j < rank – the canonical prefix form (docs/uniform_masks_vs_ranks.md). The rank masks carry the stack, so each stack element may declare its own ranks (the variety – see docs/uniform_ranks_and_varieties.md). The physical shape is a separate static int tuple (not a mask), held alongside these in .data slot 2.

Masks are STRUCTURE, so this builder emits numpy (host) arrays via np regardless of whether the supercores are jax – the np.* here is intentional and jit-required (a jax mask becomes a tracer under jit and breaks the layer). See docs/contributor/uniform_pytree_composition.md.

Parameters:
Return type:

t3toolbox.backend.common.typ.Tuple[NDArray, NDArray]