Tucker Tensor

Tucker Tensor¤

Represents the weights in a Tucker decomposition format, factorized into a core tensor and mode-specific factor matrices.

`neojax.tensor.tucker_tensor.TuckerTensor` ¤

Tucker-factorized spectral tensor.

Decomposes the spectral weight tensor \(W\) into a core tensor \(G\) and factor matrices \(U^{(j)}\) for each dimension.

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+2}} G_{r_1, \dots, r_{d+2}} U^{(1)}_{c_{out}, r_1} U^{(2)}_{c_{in}, r_2} \prod_{j=1}^d U^{(j+2)}_{x_j, r_{j+2}} \]

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+1}} G_{r_1, \dots, r_{d+1}} U^{(1)}_{c, r_1} \prod_{j=1}^d U^{(j+1)}_{x_j, r_{j+1}} \]

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 ranks, e.g. in 2D 4 ranks (out_channel, in_channel, 2 spatial ranks).	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

core_tensors (tuple[Complex[Array, ...], ...]): Learnable complex core weight tensors.
factor_matrices (tuple[Complex[Array, ...], ...]): Learnable factor matrices for each dimension.
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 tensors stacked along first axis.

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

Performs tucker tensor contraction.