ut3_apply_masks

t3toolbox.backend.ut3_masking.ut3_apply_masks(x)

def ut3_apply_masks(
        x: typ.Tuple[
            NDArray,             # tucker_supercore, shape=(d,)+stack_shape+(n,N)
            NDArray,             # tt_supercore,     shape=(d,)+stack_shape+(r,n,r)
            typ.Sequence[int],   # shape = (N0,...,N(d-1)), static int tuple
            typ.Tuple[
                NDArray,  # tucker_edge_mask, HOST bool, static
                NDArray,  # tt_edge_mask,     HOST bool, static
            ],
        ],
) -> typ.Tuple[
    NDArray,  # masked_tucker_supercore, shape=(d,)+stack_shape+(n,N)
    NDArray,  # masked_tt_supercore,     shape=(d,)+stack_shape+(r,n,r)
]:

Zero the padded (“garbage”) regions of the supercores by multiplying through the edge masks.

The mask chokepoint: every mask-using op masks on entry, so the guard here covers the whole layer. xnp.einsum on the supercore with the numpy mask as a constant operand is fine (jax promotes it). The physical shape_mask is reconstructed on the host from the static shape ints (np, never jnp – a traced mask breaks the layer; see docs/contributor/uniform_pytree_composition.md).

Parameters:

x (t3toolbox.backend.common.typ.Tuple[NDArray, NDArray, t3toolbox.backend.common.typ.Sequence[int], t3toolbox.backend.common.typ.Tuple[NDArray, NDArray]])

Return type:

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