Dense Tensor

Standard Dense Tensor

A standard un-factorized representation of Fourier weights.

neojax.tensor.dense_tensor.DenseTensor

Standard dense spectral tensor.

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
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

• weights (tuple[Complex[Array, ...], ...]): Learnable complex weights for the retained Fourier modes.
• separable (bool): Whether to use separable implementation of contraction.

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

Reconstruct dense tensors into single dense tensor.

Returns:

Type	Description
Complex[Array, 'n_corners ...']	Dense tensors stacked along first axis.

from_weights(weights: Complex[Array, 'n_corners ...'], separable: bool = False) -> DenseTensor classmethod

Alternative constructor to wrap constructed dense tensors.

Commonly used to construct a DenseTensor from factorized tensors after calling to_dense(), e.g. if reconstructed flag is True.

Parameters:

Name	Type	Description	Default
weights	Complex[Array, 'n_corners ...']	Constructed dense tensors.	required
separable	bool	Whether the wrapped weights are separable.	False

Returns:

Type	Description
DenseTensor	A new instance of DenseTensor.

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

Performs dense tensor contraction.

Parameters:

Name	Type	Description	Default
corner_idx	int	Corner index of the n-dim FFT hypercube.	required
x_ft_slice	Complex[Array, 'in_c ...']	Mode slice of fourier transformed input.	required

Returns:

Type	Description
Complex[Array, ...]	Tensor contraction of respective weight and input.