General

General

In Neural Operators (such as the FNO and TFNO), weights are parameterized in the Fourier domain. For high-dimensional or high-channel settings, storing the full parameter tensor \((C_{out}, C_{in}, m_1, \dots, m_d)\) is computationally expensive. neojax provides modular tensor classes to represent these weights in dense or low-rank factorized forms.

Useful overview on Tensor Operations

Tensor Decompositions and Applications

@article{kolda2009tensor,
    title={Tensor decompositions and applications},
    author={Kolda, Tamara G and Bader, Brett W},
    journal={SIAM review},
    volume={51},
    number={3},
    pages={455--500},
    year={2009},
    publisher={SIAM}
}

---

Separable Weights (separable=True)

All tensor classes support depthwise-separable spectral convolutions via the separable=True flag.

When separable=True:

• Input and output channels must match (in_channels == out_channels).
• The weight tensor dimensions are reduced from \((C_{out}, C_{in}, m_1, \dots, m_d)\) to \((C, m_1, \dots, m_d)\).
• Contraction behaves pointwise/elementwise across the channel dimension. This is the spectral analogue of a depthwise separable convolution, which lowers parameter counts and memory usage.