neuralop.models.TFNO

class neuralop.models.TFNO(*args, **kwargs)[source]

Tucker Tensorized Fourier Neural Operator (TFNO).

TFNO is an FNO with Tucker factorization enabled by default.

It uses Tucker factorization of the weights, making the forward pass efficient by contracting directly with the factors of the decomposition.

This results in a fraction of the parameters of an equivalent dense FNO.

Parameters:

factorizationstr, optional

Tensor factorization method, by default “Tucker”

rankfloat, optional

Tensor rank for factorization, by default 0.1. A TFNO with rank 0.1 has roughly 10% of the parameters of a dense FNO.

All other parameters are inherited from FNO with identical defaults.

See FNO class docstring for the complete parameter list.

Examples

>>> from neuralop.models import TFNO
>>> # Create a TFNO model with default Tucker factorization
>>> model = TFNO(n_modes=(12, 12), in_channels=1, out_channels=1, hidden_channels=64)
>>>
>>> # Equivalent FNO model with explicit factorization:
>>> model = FNO(n_modes=(12, 12), in_channels=1, out_channels=1, hidden_channels=64,
...             factorization="Tucker", rank=0.1)