CP Tensor

CP (Canonical-Polyadic) Tensor¤

Represents the weights in a Canonical Polyadic decomposition format (approximating the weight tensor as a sum of K rank-1 tensors).

`neojax.tensor.cp_tensor.CPTensor` ¤

Canonical-Polyadic-decomposed Tensor.

Decomposes the spectral weight tensor \(W\) into a sum of \(R\) rank-1 tensors (where \(R\) is the CP rank), represented by a core vector \(\lambda\) and factor matrices \(U^{(j)}\) for each mode.

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}^R \lambda_r U^{(1)}_{c_{out}, r} U^{(2)}_{c_{in}, r} \prod_{j=1}^d U^{(j+2)}_{x_j, r} \]

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}^R \lambda_r U^{(1)}_{c, r} \prod_{j=1}^d U^{(j+1)}_{x_j, r} \]

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`	Single rank for CP tensor. In a CP tensor, all modes share the same rank.	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, "r"], ...]): 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): Spatial dimensions of the tensor structure.
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 cp tensor contraction.