ufv_apply_frame_masks

t3toolbox.backend.ufv_masking.ufv_apply_frame_masks(data)

def ufv_apply_frame_masks(
        data: typ.Tuple[
            NDArray,             # up_tucker_supercore,  (d,)+stack_shape+(nU, N)
            NDArray,             # down_tt_supercore,    (d,)+stack_shape+(rL, nD, rR)
            NDArray,             # left_tt_supercore,    (d,)+stack_shape+(rL, nU, rL)
            NDArray,             # right_tt_supercore,   (d,)+stack_shape+(rR, nU, rR)
            typ.Sequence[int],   # shape = (N0,...,N(d-1)), static int tuple
            typ.Tuple[
                NDArray,  # up_mask,          dtype=bool, (d,)  +stack_shape+(nU,)
                NDArray,  # down_mask,        dtype=bool, (d,)  +stack_shape+(nD,)
                NDArray,  # frame_left_mask,  dtype=bool, (d+1,)+stack_shape+(rL,)
                NDArray,  # frame_right_mask, dtype=bool, (d+1,)+stack_shape+(rR,)
            ],
        ],
) -> typ.Tuple[
    NDArray,  # masked_up_tucker_supercore
    NDArray,  # masked_down_tt_supercore
    NDArray,  # masked_left_tt_supercore
    NDArray,  # masked_right_tt_supercore
]:

Zero the padded (“garbage”) regions of the frame supercores via the edge masks. 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:

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

Return type:

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