Tensor Train Tensor

Tensor Train (TT) Tensor¤

Represents the weights in a Tensor Train format, factorizing the tensor into a chain of low-dimensional tensors connected back-to-back.

`neojax.tensor.tt_tensor.TTTensor` ¤

Tensor Train Tensor.

Decomposes the spectral weight tensor \(W\) into a train/chain of low-dimensional tensors \(U^{(j)}\) connected back-to-back.

For the standard case (separable=False), the weight tensor \(W \in \mathbb{C}^{C_{out} \times C_{in} \times m_1 \times \dots \times m_d}\) is reconstructed as:

\[ W_{c_{out}, c_{in}, x_1, \dots, x_d} = \sum_{r_1, \dots, r_{d+1}} U^{(1)}_{c_{out}, r_1} U^{(2)}_{r_1, c_{in}, r_2} U^{(3)}_{r_2, x_1, r_3} \dots U^{(d+2)}_{r_{d+1}, x_d} \]

For the separable case (separable=True), the weight tensor \(W \in \mathbb{C}^{C \times m_1 \times \dots \times m_d}\) is reconstructed as:

\[ W_{c, x_1, \dots, x_d} = \sum_{r_1, \dots, r_d} U^{(1)}_{c, r_1} U^{(2)}_{r_1, x_1, r_2} \dots U^{(d+1)}_{r_d, x_d} \]

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for weight initialization.	required
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int`	Number of output channels.	required
`modes`	`Sequence[int]`	Number of Fourier modes to retain across each spatial dim.	required
`ranks`	`int \| Sequence[int]`	Number of ranks to contract the spectral tensors to. If `ranks` is an integer, the same number is used for all ranks. If not, should be (num_spatial_dims + num_channel_dims) - 1 ranks, e.g. in 2D 3 ranks (out_channel, in_channel, 2 spatial ranks) - 1.	required
`init_std`	`Union[float, Literal['auto']]`	Standard deviation to use for weight initialization, by default 'auto'. If 'auto', uses (2 / (in_channels + out_channels)) ** 0.5.	`'auto'`
`separable`	`bool`	Whether to use separable implementation of contraction. If True, contracts factors of factorized tensor weight individually. Default is False.	`False`

Internal Attributes

lr_tensors (tuple[Complex[Array, "..."], ...]): Learnable complex low-rank tensors.
einsum_str (str): Contraction string passed directly to jnp.einsum.
to_dense_einsum_str (str): Contraction string passed to jnp.einsum in to_dense.
ndim (int): Total dimensions of the tensor structure (channels + spatial).
separable (bool): Whether to use separable implementation of contraction.

`to_dense() -> Complex[Array, 'n_corners ...']` ¤

Reconstruct factorized tensors into single dense tensor.

Useful, e.g. if the reconstructed flag is True.

Returns:

Type	Description
`Complex[Array, 'n_corners ...']`	Reconstructed dense tensor stacked along first axis.

`call(corner_idx: int, x_ft_slice: Complex[Array, 'in_c ...']) -> Complex[Array, ...]` ¤

Performs tensor train tensor contraction.