Skip to content

FNO Blocks

FNO Blocks¤

The standard layer of the Fourier Neural Operator. An FNO Block computes the sum of the global spectral convolution and a local skip connection, followed by optional normalization and non-linear activation.

neojax.nn.fno_blocks.FNOBlocks ¤

General FNO Blocks with variable layer number.

Implemented as in [1] and [2]. Each block is an instance of FNOBlock with an optional channelwise MLP.

Parameters:

Name Type Description Default
key PRNGKeyArray

PRNG key for parameter initialization.

 required
n_layers int

Number of consecutive FNO blocks.

 required
in_channels int

Number of input channels.

 required
out_channels int

Number of output channels.

 required
modes int | Sequence[int]

Number of Fourier modes to retain across each spatial dimension, e.g. (16, 16) for 2D inputs.

 required
activation Callable | Sequence[Callable]

Activation function or sequence of activation functions used inside the FNO blocks. Default is jax.nn.gelu.

 gelu
use_channel_mlp bool

Whether to apply a pointwise channel MLP with 2 layers after each FNO block. Default is True.

 True
preactivation bool

Whether to apply the activation function before the spectral convolution and skip connections. Default is False. This is an exclusive flag. If activation is applied before, it isn't applied after.

 False
normalization Optional[Literal['layer', 'instance', 'group']]

Type of normalization to use. Applied after the spectral_op + local_op sum, before activation. Can be "layer", "instance", "group" or None. Default is "layer".

 'layer'
norm_groups int

Number of groups to use if normalization="group". Defaults to 1.

 1
use_fno_residual bool

Whether to use a Resnet-style residual connection around each FNO block. Improves stability. Default True.

 True
local_operator Optional[Literal['linear', 'soft-gating', 'identity']]

Type of local operator to use in the FNO block. Can be "linear", "soft-gating", "identity", or None. Default is "linear".

 'linear'
use_local_operator_bias bool

Whether to use a bias term in the FNO blocks local operator. Default is False.

 False
channel_mlp_residual Optional[Literal['linear', 'soft-gating', 'identity']]

Type of residual connection used around the channel MLPs. Can be "linear", "soft-gating", "identity", or None. Default is "identity", which corresponds to Resnet-style residual connection.

 'identity'
channel_mlp_expansion float | None

Expansion factor for computing the hidden channel dimension of the MLPs. Default is 0.5.

 0.5
channel_mlp_activations Callable | Sequence[Callable]

Activation function or sequence of activation functions used inside the channel MLPs. Default is jax.nn.gelu.

 gelu
channel_mlp_dropout float

Dropout probability applied after each layer (except the last) of the channel MLPs. Defaults to 0.0.

 0.0
enforce_hermitian_symmetry bool

Whether to enforce hermitian symmetry on the outputs of the spectral convolutions before calling irfftn. Default is True. If set to False, cuFFT on GPU may cause line artifacts when calling irfftn.

 True
fft_norm Optional[Literal['forward', 'backward', 'ortho']]

FFT normalization. Can be None, "backward", "ortho" or "forward". Default is "forward".

 'forward'
is_complex_data bool

Whether input data is complex valued. If True, uses full FFT. Default is False.

 False
resolution_scaling_factor float | int | Sequence[float | int] | None

Layerwise factor by which to scale the domain resolution of the function. Default is None. Passing a scalar scales the resolution by that value each layer. Passing a sequence, scales the i-th layer by the i-th value.

 None
ranks int | Sequence[int] | None

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 a sequence of ranks. Default is None.

 None
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'
factorization Union[BaseTensor, Literal['tucker', 'cp', 'tt'], NoneType]

Tensor factorization type. Can be a BaseTensor instance, a string ("tucker", "cp", "tt"), or None. Default is None.

 None
implementation Literal['reconstructed', 'factorized']

Weight reconstruction mode. Can be "reconstructed" or "factorized". Default is "factorized".

 'factorized'
separable bool

Whether to use separable implementation of contraction. If True, contracts factors of factorized tensor weight individually. Default is False.

 False
Internal Attributes

These fields store the internal layers state (and weights).

  • fno_layers (tuple[FNOBlock, ...]): The initialized FNOBlock layers.
  • channel_mlps (tuple[PointwiseMLP, ...] | None): The initialized PointwiseMLP layers, or None if use_channel_mlp is False.
  • channel_mlp_residuals (tuple | None): The residual connection instances or None.

Examples:

import jax.random as jr
import jax.numpy as jnp
from neojax.nn.fno_blocks import FNOBlocks

key = jr.key(0)

# Initialize a sequence of 4 FNO Blocks
fno_blocks = FNOBlocks(
    key=key,
    n_layers=4,
    in_channels=3,
    out_channels=16,
    modes=(8, 8),
    use_channel_mlp=True
)
# Input shape: (channels, height, width)
x = jnp.ones((3, 64, 64))
out = fno_blocks(x)
Cite

Fourier Neural Operator for Parametric Partial Differential Equations

@inproceedings{
    li2021fourier,
    title={Fourier Neural Operator for Parametric Partial Differential Equations},
    author={Zongyi Li and Nikola Kovachki and
    Kamyar Azizzadenesheli and Burigede liu and
    Kaushik Bhattacharya and Andrew Stuart and Anima Anandkumar},
    booktitle={International Conference on Learning Representations},
    year={2021},
    url={https://openreview.net/forum?id=c8P9NQVtmnO}
}

Neural Operator: Learning Maps Between Function Spaces With Applications to PDEs

@article{kovachki2023neural,
    title={Neural operator: Learning maps between function spaces with applications to pdes},
    author={Kovachki, Nikola and Li, Zongyi and
    Liu, Burigede and Azizzadenesheli,
    Kamyar and Bhattacharya, Kaushik and Stuart, Andrew and Anandkumar, Anima},
    journal={Journal of Machine Learning Research},
    volume={24},
    number={89},
    pages={1--97},
    year={2023}
}
__call__(x: Inexact[Array, 'in_c ...'], *, key: PRNGKeyArray | None = None, inference: bool = False) -> Inexact[Array, 'out_c ...'] ¤

Forward pass through n_layers of FNO Blocks.

Parameters:

Name Type Description Default
x Inexact[Array, 'in_c ...']

Input array.

 required
key PRNGKeyArray | None

PRNG key used for dropout masks.

 None
inference bool

If True, dropout is disabled.

 False

Returns:

Type Description
Inexact[Array, 'out_c ...']

Output array.

neojax.nn.fno_blocks.FNOBlock ¤

General FNO Block.

FNO block composed of the global non-linear integral kernel, i.e., \(\mathcal{K}(\mathbf{v}) = \mathcal{F}^{-1}(R \cdot \mathcal{F}(\mathbf{v}))\), and an optional local operator \(W \mathbf{v}\). A block can be further customized with Resnet-style residual connections around the FNO and normalization before the non-linear activation.

Parameters:

Name Type Description Default
key PRNGKeyArray

PRNG key for parameter initialization.

 required
in_channels int

Number of input channels.

 required
out_channels int

Number of output channels.

 required
modes int | Sequence[int]

Number of Fourier modes to retain across each spatial dimension. Must be a sequence of integers (e.g., (16, 16) for 2D) or single integer for 1D inputs.

 required
activation Callable

Activation function to use. Defaults to jax.nn.gelu.

 gelu
local_operator Optional[Literal['linear', 'soft-gating', 'identity']]

Type of local operator to use. Can be "linear", "soft-gating", "identity", or None. Defaults to "linear".

 'linear'
use_local_operator_bias bool

Whether to use a bias term in the FNO blocks local operator. Defaults to False.

 False
normalization Optional[Literal['layer', 'instance', 'group']]

Type of normalization to use. Applied after the spectral_op + local_op sum, before activation. Can be "layer", "instance", "group" or None. Default is "layer".

 'layer'
norm_groups int

Number of groups to use if normalization="group". Defaults to 1.

 1
use_fno_residual bool

Whether to use a Resnet-style residual connection around each FNO block. Improves stability. Default True.

 True
preactivation bool

Whether to apply the activation before the spectral convolution and skip connection. Defaults to False. This is an exclusive flag. If activation is applied before, it isn't applied after.

 False
enforce_hermitian_symmetry bool

Whether to enforce hermitian symmetry on the outputs of the spectral convolutions before calling irfftn. Default is True. If set to False, the outputs will not be hermitian-symmetric and the output will not be real-valued. Only considered for real-valued inputs.

 True
fft_norm Optional[Literal['forward', 'backward', 'ortho']]

FFT normalization. Can be None, "backward", "ortho" or "forward". Default is "forward".

 'forward'
is_complex_data bool

Whether input data is complex valued. If True, uses full FFT. Default is False.

 False
resolution_scaling_factor float | int | None

Factor by which to scale the domain resolution of the function. Default is None, no scaling.

 None
ranks int | Sequence[int] | None

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 a sequence of ranks. Default is None.

 None
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'
factorization Union[BaseTensor, Literal['tucker', 'cp', 'tt'], NoneType]

Tensor factorization type. Can be a BaseTensor instance, a string ("tucker", "cp", "tt"), or None. Default is None.

 None
implementation Literal['reconstructed', 'factorized']

Weight reconstruction mode. Can be "reconstructed" or "factorized". Default is "factorized".

 'factorized'
separable bool

Whether to use separable implementation of contraction. If True, contracts factors of factorized tensor weight individually. Default is False.

 False
Internal Attributes

These fields store the internal layers state (and weights).

  • spectral_conv (SpectralConvNd): The SpectralConvNd layer performing the operator integral.
  • local_operator (Flattened1dConv | SoftGating | eqx.nn.Identity | None): The initialized local operator or None.
  • normalization (InstanceNorm | eqx.nn.GroupNorm | None): Type of normalization to use. Applied after the spectral_op + local_op sum, before activation.
  • activation (Callable): The activation function.
  • preactivation (bool): Boolean flag indicating if preactivation is used.
  • use_fno_residual (bool): Whether to use a Resnet-style residual connection around each FNO block. Improves stability.

Examples:

import jax.random as jr
from neojax.nn.fno_blocks import FNOBlock
import jax.numpy as jnp

key = jr.key(0)
# Initialize a 2D FNO Block
fno_block = FNOBlock(
    key=key,
    in_channels=3,
    out_channels=8,
    modes=(4, 4),
    local_operator="linear"
)
# Input shape: (channels, height, width)
x = jnp.ones((3, 32, 32))
out = fno_block(x)
__call__(x: Inexact[Array, 'in_c ...']) -> Inexact[Array, 'out_c ...'] ¤

Forward pass of FNO block.

Parameters:

Name Type Description Default
x Inexact[Array, 'in_c ...']

Input array.

 required

Returns:

Type Description
Inexact[Array, 'out_c ...']

Output array.