ut3_to_dense

t3toolbox.backend.ut3_conversions.ut3_to_dense(x)

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

Form the dense tensor from a uniform Tucker tensor train: mask the supercores, chain-contract (shared with the ragged path), then static-prefix-slice the padded physical axes to the real shape.

The real shape is the static int tuple x[2]; the padded physical axes are sliced to it.

Jitting this functionally (no frontend): close over the host masks as constants and trace only the supercores. Tracing the whole .data instead would put the masks among the traced args, which the guard rejects (docs/contributor/uniform_pytree_composition.md):

>>> import numpy as np, jax
>>> import t3toolbox.tucker_tensor_train as t3
>>> import t3toolbox.uniform_tucker_tensor_train as ut3
>>> import t3toolbox.backend.ut3_conversions as ut3_conversions
>>> np.random.seed(0)
>>> ux = ut3.UniformTuckerTensorTrain.from_t3(t3.TuckerTensorTrain.randn((4, 5, 6), (2, 3, 2), (1, 2, 2, 1))).to_jax()
>>> tk, tt, shape, masks = ux.data                          # shape ints + HOST bool masks, static
>>> dense = jax.jit(lambda a, b: ut3_conversions.ut3_to_dense((a, b, shape, masks)))(tk, tt)  # RIGHT
>>> bool(np.allclose(dense, ux.to_dense()))
True
>>> jax.jit(ut3_conversions.ut3_to_dense)(ux.data)  # WRONG: masks are traced args
Traceback (most recent call last):
ValueError: uniform masks must be concrete host (numpy) arrays, ...

Parameters:

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

Return type:

NDArray