Models Reference¤

This page contains the API reference for all pre-built models in neojax. Currently, the library supports the Fourier Neural Operator (FNO), Tucker-factorized FNO (TFNO), and Deep Operator Networks (DeepONet).

Improved Parameter Naming (FNO & TFNO)¤

If you are migrating from the original PyTorch neuraloperator library, you will notice that neojax introduces different parameter naming convention for residual connections and skip connections. The original library relies on ambiguous parameter names (like fno_skip).

To provide a clean, modern API, neojax enforces the following standard across all FNO and TFNO architectures:

Any parameters containing the word ...local_operator... refer to the local operator.
Any parameters containing the word ...residual... refer to Resnet-style residual connections around components.

Fourier Neural Operator (FNO)¤

`neojax.models.fno.FNO` ¤

General n-dimensional Fourier Neural Operator (FNO).

The model consists of a lifting layer that maps the input to a higher-dimensional latent space, a sequence of FNO blocks (spectral convolutions + skip connections) optionally interleaved with pointwise MLPs, and a final projection layer.

The implementation is in accordance with [1] and [2].

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for parameter initialization.	required
`in_channels`	`int`	Number of input channels (e.g., coordinates + initial conditions).	required
`out_channels`	`int`	Number of output channels (e.g., solution field).	required
`hidden_channels`	`int`	Hidden channel dimension (latent width) used throughout the FNO blocks. This significantly affects the number of parameters. Good starting point is 64 and increase if needed. Update `lift_channel_ratio` and `proj_channel_ratio` accordingly. They scale proportional to hidden_channels.	required
`n_layers`	`int`	Number of consecutive FNO blocks.	required
`modes`	`int \| Sequence[int]`	Number of Fourier modes to retain across each spatial dimension.	required
`activation`	`Callable`	Activation function used within the layers. Defaults to `jax.nn.gelu`.	`gelu`
`use_channel_mlp`	`bool`	Whether to apply a pointwise channel MLP after each FNO block. Defaults to `True`.	`True`
`local_operator`	`Literal['linear', 'soft-gating', 'identity'] \| None`	Type of skip connection inside FNO blocks. 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. Default is `False`.	`False`
`channel_mlp_residual`	`Literal['linear', 'soft-gating', 'identity'] \| None`	Type of skip connection around channel MLPs. Can be `"linear"`, `"soft-gating"`, `"identity"`, or None. Defaults to `"soft-gating"`.	`'soft-gating'`
`channel_mlp_expansion`	`float \| None`	Expansion factor for hidden dimension in the channel MLPs. Defaults to 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`
`normalization`	`Literal['layer', 'instance', 'group'] \| None`	Type of normalization to use in FNO blocks. Can be `"layer"`, `"instance"`, `"group"` or None. Default is `"layer"`.	`'layer'`
`norm_groups`	`int`	Number of groups for group normalization. Default is 1.	`1`
`use_fno_residual`	`bool`	Whether to use residual connection around FNO blocks. Default is True.	`True`
`preactivation`	`bool`	Whether to use pre-activation style blocks. Default is False.	`False`
`n_lift_layers`	`int`	Number of layers in the lifting MLP. Defaults to 2.	`2`
`n_proj_layers`	`int`	Number of layers in the projection MLP. Defaults to 2.	`2`
`lift_channel_ratio`	`float`	Ratio of lifting channels to hidden_channels. The number of lifting channels in the lifting block of the FNO is lifting_channel_ratio * hidden_channels (e.g. default 2 * hidden_channels).	`2.0`
`proj_channel_ratio`	`float`	Ratio of projection channels to hidden_channels. The number of projection channels in the projection block of the FNO is projection_channel_ratio * hidden_channels (e.g. default 2 * hidden_channels).	`2.0`
`positional_embedding`	`GridEmbeddingNd \| None`	Positional embedding to apply to last channels of raw input before passing through FNO. Defaults to regular GridEmbeddingNd on a ((0, 1), ...) grid.	`None`
`domain_padding`	`float \| Sequence[float] \| None`	Percentage of padding to use. If single float, this padding is used for all dims. Sequence of floats indicates padding percentage per dim. Default is None, no padding.	`None`

Internal Attributes

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

positional_embedding (GridEmbeddingNd | None): Positional embedding to apply to last channels of raw input before passing through FNO.
lifting (PointwiseMLP): The PointwiseMLP used to lift inputs to the hidden hidden_channels.
fno_blocks (FNOBlocks): The FNOBlocks sequence containing spectral convolutions.
projection (PointwiseMLP): The PointwiseMLP used to project latent features to out_channels.
padding (DomainPadding | None): Percentage of domain padding to use.

Cite

[Fourier Neural Operator for Parametric Partial Differential Equations] (https://arxiv.org/abs/2010.08895)

@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] (https://www.jmlr.org/papers/volume24/21-1524/21-1524.pdf)

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

Upcoming Features

Current implementation doesn't support different FNO block precisions, resolution scaling, stabilizers, separable spectral convolutions or enforcing hermitian symmetry. These will be added in future releases.

Tucker-factorized FNO (TFNO)¤

`neojax.models.tfno.TFNO` ¤

General n-dimensional Tucker factorized FNO.

The model consists of a lifting layer that maps the input to a higher-dimensional latent space, a sequence of TFNO blocks (spectral convolutions + skip connections) optionally interleaved with pointwise MLPs, and a final projection layer.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for parameter initialization.	required
`in_channels`	`int`	Number of input channels (e.g., coordinates + initial conditions).	required
`out_channels`	`int`	Number of output channels (e.g., solution field).	required
`hidden_channels`	`int`	Hidden channel dimension (latent width) used throughout the TFNO blocks. This significantly affects the number of parameters. Good starting point is 64 and increase if needed. Update `lift_channel_ratio` and `proj_channel_ratio` accordingly. They scale proportional to hidden_channels.	required
`n_layers`	`int`	Number of consecutive TFNO blocks.	required
`modes`	`int \| Sequence[int]`	Number of Fourier modes to retain across each spatial dimension.	required
`ranks`	`int \| Sequence[int]`	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 num_spatial_dims + num_channel_dims ranks, e.g. in 2D 4 ranks (out_channel, in_channel, 2 spatial ranks).	required
`share_factor_matrices`	`bool`	Whether to share the factor matrices across the weight tensors. This further decreases the number of weight tensors. Only the core tensor is not shared then. Default is True.	`True`
`activation`	`Callable`	Activation function used within the layers. Defaults to `jax.nn.gelu`.	`gelu`
`use_channel_mlp`	`bool`	Whether to apply a pointwise channel MLP after each TFNO block. Defaults to `True`.	`True`
`local_operator`	`Literal['linear', 'soft-gating', 'identity'] \| None`	Type of skip connection inside TFNO blocks. 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. Default is `False`.	`False`
`channel_mlp_residual`	`Literal['linear', 'soft-gating', 'identity'] \| None`	Type of skip connection around channel MLPs. Can be `"linear"`, `"soft-gating"`, `"identity"`, or None. Defaults to `"soft-gating"`.	`'soft-gating'`
`channel_mlp_expansion`	`float \| None`	Expansion factor for hidden dimension in the channel MLPs. Defaults to 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`
`normalization`	`Literal['layer', 'instance', 'group'] \| None`	Type of normalization to use in TFNO blocks. Can be `"layer"`, `"instance"`, `"group"` or None. Default is `"layer"`.	`'layer'`
`norm_groups`	`int`	Number of groups for group normalization. Default is 1.	`1`
`use_fno_residual`	`bool`	Whether to use residual connection around TFNO blocks. Default is True.	`True`
`preactivation`	`bool`	Whether to use pre-activation style blocks. Default is False.	`False`
`n_lift_layers`	`int`	Number of layers in the lifting MLP. Defaults to 2.	`2`
`n_proj_layers`	`int`	Number of layers in the projection MLP. Defaults to 2.	`2`
`lift_channel_ratio`	`float`	Ratio of lifting channels to hidden_channels. The number of lifting channels in the lifting block of the TFNO is lifting_channel_ratio * hidden_channels (e.g. default 2 * hidden_channels).	`2.0`
`proj_channel_ratio`	`float`	Ratio of projection channels to hidden_channels. The number of projection channels in the projection block of the TFNO is projection_channel_ratio * hidden_channels (e.g. default 2 * hidden_channels).	`2.0`
`positional_embedding`	`GridEmbeddingNd \| None`	Positional embedding to apply to last channels of raw input before passing through TFNO. Defaults to regular GridEmbeddingNd on a ((0, 1), ...) grid.	`None`
`domain_padding`	`float \| Sequence[float] \| None`	Percentage of padding to use. If single float, this padding is used for all dims. Sequence of floats indicates padding percentage per dim. Default is None, no padding.	`None`

Internal Attributes

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

positional_embedding (GridEmbeddingNd | None): Positional embedding to apply to last channels of raw input before passing through TFNO.
lifting (PointwiseMLP): The PointwiseMLP used to lift inputs to the hidden hidden_channels.
tfno_blocks (TFNOBlocks): The TFNOBlocks sequence containing spectral convolutions.
projection (PointwiseMLP): The PointwiseMLP used to project latent features to out_channels.
padding (DomainPadding | None): Percentage of domain padding to use.

Info

The TFNO has the same architecture as the FNO except for the factorized spectral weights.

Upcoming Features

Current implementation doesn't support different TFNO block precisions, resolution scaling, stabilizers, separable spectral convolutions or enforcing hermitian symmetry. These will be added in future releases.

Deep Operator Networks (DeepONet)¤

neojax implements a general modular DeepONet (arbitrary branch and trunk networks) and a preconfigured MLPDeepONet (MLP branch and MLP trunk network).

The DeepONets are implemented to evaluate the underlying operator with input function \(u \in \mathbb{R}^m\) at the output \(y \in \mathbb{R}^d\). In case of higher dimensional input spaces, e.g. \(\mathbb{R}^2, \mathbb{R}^3, \dots\) the sensor locations m may not be 1D arrays, as long as the branch network handles this correctly.

In practice, it is more useful to evaluate the output across a grid of points y. See below for details.

How to evaluate the input function across a grid of points: To evaluate a single input function \(u\) across a grid (or batch) of points \(y\), use jax.vmap. Here is how you would evaluate on a \(256 \times 256\) grid:

import jax
import jax.numpy as jnp
from neojax.models import DeepONet

# Initialize model
model = DeepONet(...)

# Define inputs
u = jnp.ones((100,))  # One function input
y_grid = jnp.meshgrid(jnp.linspace(0, 1, 256), jnp.linspace(0, 1, 256))
y_points = jnp.stack(y_grid, axis=-1)  # Shape: (256, 256, 2)

# Vectorize the model over the coordinate axes
# in_axes: (None, 0) means 'u' is fixed, 'y' is mapped over its 0-th dimension
vmapped_inner = jax.vmap(model, in_axes=(None, 0))          # Maps (256, 2) -> (256, 1)
vmapped_outer = jax.vmap(vmapped_inner, in_axes=(None, 0))  # Maps (256, 256, 2) -> (256, 256, 1)

# Generate predictions
predictions = vmapped_outer(u, y_points)  # Shape: (256, 256, 1)

Pro Tip: If you want to evaluate a batch of functions across a batch of points, just add a third jax.vmap call.

# Batch of 32 functions, each evaluated at 256 points
u_batch = jnp.ones((32, 100))
y_batch = jnp.zeros((32, 256, 2))

# vmap over functions (axis 0) and points (axis 0)
batch_model = jax.vmap(vmapped_inner, in_axes=(0, 0))
batch_preds = batch_model(u_batch, y_batch) # Shape: (32, 256, 1)

General DeepONet¤

`neojax.models.deeponet.DeepONet` ¤

General DeepONet (Deep Operator Network).

Learns a PDE operator through two separate branch and trunk networks. The branch net encodes the discrete function space and the trunk net encodes the domain of output functions. Implemented as in the original publication [1].

Parameters:

Name	Type	Description	Default
`branch_net`	`Module`	The network representing the branch (encodes u). Should map shape (m_sensors,) to (p,). May also take different input shapes, as long as it outputs (p,), e.g. CNN branch net.	required
`trunk_net`	`Module`	The network representing the trunk (encodes y). Should map shape (d_dim,) to (p,).	required
`out_activation`	`Callable \| None`	Optional output activation after dot product.	`None`
`bias`	`Float[Array, 1] \| None`	Optional learnable bias after dot product.	`None`

Internal Attributes

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

branch_net (eqx.Module): Arbitrary NN that maps (m_sensors,) -> (p,).
trunk_net (eqx.Module): Arbitrary NN that maps (d_dim,) -> (p,).
out_activation (Callable | None): Optional output activation after dot product.
bias (Float[Array, "1"] | None): Optional learnable bias after dot product.

Info

This is the Unstacked (Standard) implementation as described in Lu et al. (2021). It uses a single branch network and a single trunk network for maximum computational efficiency.

For vector-valued operators (the "Stacked" variant), it is recommended to either:

Increase the output dimension of the branch/trunk networks.
Use jax.vmap to wrap this class for multiple independent operators.

Cite

[Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators] (https://www.nature.com/articles/s42256-021-00302-5)

@article{lu2021learning,
    title={Learning nonlinear operators via DeepONet
    based on the universal approximation theorem of operators},
    author={Lu, Lu and Jin, Pengzhan and Pang,
    Guofei and Zhang, Zhongqiang and Karniadakis, George Em},
    journal={Nature machine intelligence},
    volume={3},
    number={3},
    pages={218--229},
    year={2021},
    publisher={Nature Publishing Group}
}

MLPDeepONet¤

`neojax.models.deeponet.MLPDeepONet` ¤

DeepONet with MLP branch and MLP trunk network.

The branch MLP maps (m_sensors,) funtion values to the (p,) latent vector and the trunk MLP maps the (d,) evaluation vector to the (p,) latent vector.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	Jax random key.	required
`m_sensors`	`int`	Number of function values u.	required
`d_dim`	`int`	Output space dimension.	required
`p_latent`	`int`	Latent vector dimension.	required
`branch_hidden_dims`	`Sequence[int]`	Hidden dimensions of the branch MLP. First layer is (m_sensors, branch_hidden_dims[0]) and last layer is (branch_hidden_dims[-1], p_latent).	required
`trunk_hidden_dims`	`Sequence[int]`	Hidden dimensions of the trunk MLP. First layers is (d_dim, trunk_hidden_dims[0]) and last layer is (trunk_hidden_dims[-1], p_latent).	required
`branch_activations`	`Callable \| Sequence[Callable]`	Activation functions of branch MLP.	`<function gelu>`
`trunk_activations`	`Callable \| Sequence[Callable]`	Activation functions of trunk MLP. Single callable means, all activations use this function.	`<function gelu>`
`out_activation`	`Callable \| None`	Optional output activation after dot product.	`None`
`bias`	`Float[Array, '1'] \| None`	Optional learnable bias after dot product.	`None`

Internal Attributes

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

branch_net (PointwiseMLP): MLP that maps (m_sensors,) -> (p,).
trunk_net (PointwiseMLP): MLP that maps (d_dim,) -> (p,).
out_activation (Callable | None): Optional output activation after dot product.
bias (Float[Array, "1"] | None): Optional learnable bias after dot product.

Models Reference¤

Improved Parameter Naming (FNO & TFNO)¤

Fourier Neural Operator (FNO)¤

neojax.models.fno.FNO ¤

Tucker-factorized FNO (TFNO)¤

neojax.models.tfno.TFNO ¤

Deep Operator Networks (DeepONet)¤

General DeepONet¤

neojax.models.deeponet.DeepONet ¤

MLPDeepONet¤

neojax.models.deeponet.MLPDeepONet ¤

`neojax.models.fno.FNO` ¤

`neojax.models.tfno.TFNO` ¤

`neojax.models.deeponet.DeepONet` ¤

`neojax.models.deeponet.MLPDeepONet` ¤