General

General

Shared Model Functionality

All models inherit useful general functionality from BaseNO such as:

• size(): Get the model size in MB.
• astype(): Cast model weights to a different datatype.
• profile_compile(): Try to jit-compile the model and get a compilation summary.
• load(): Load model and hyperparameters.
• save(): Save model and hyperparameters.
• load_weights(): Load only the model weights.
• save_weights(): Save only the model weights.

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.

Stateful Layers

Some layers (e.g. batch normalization, as in CNO) carry states, like running statistics. Following Equinox, such models are created with eqx.nn.make_with_state and are called as out, state = model(x, state), inside of a jax.vmap with a matching axis_name. model.is_stateful() tells whether a model (or layer) is stateful. Models can use apply_layer to call layers that may or may not be stateful uniformly:

neojax.models.utils.apply_layer(layer: Module, x: PyTree, state: equinox.nn._stateful.State | None = None, **kwargs: typing.Any) -> tuple[PyTree, equinox.nn._stateful.State | None]

Applies a layer that may or may not be stateful.

Follows the convention of eqx.nn.Sequential: A layer is called as out, state = layer(x, state) if it is an eqx.nn.StatefulLayer with layer.is_stateful(), otherwise as out = layer(x). This allows to write models in terms of a threaded state that does not care whether individual layers (or the whole model, depending on its configuration) have state, see Equinox docs.

Parameters:

Name	Type	Description	Default
layer	Module	The layer to call.	required
x	PyTree	Input of the layer.	required
state	State | None	The current state or None if there is none.	None
kwargs	Any	Keyword arguments (e.g. key or inference) passed on to the layer in both cases.	required

Returns:

Type	Description
tuple[PyTree, State | None]	The output of the layer and the (updated) state. A stateless layer returns the state unchanged.

Examples:

import equinox as eqx
import jax.numpy as jnp
import jax.random as jr
from neojax.models.utils import apply_layer

key = jr.key(0)
x = jnp.ones((3, 8))

# Stateless layer: state (here `None`) is passed through
linear = eqx.nn.Linear(8, 4, key=key)
out, state = apply_layer(linear, x, None)

Stochastic Layers & PRNG Key Management

In neojax, models that use stochastic layers (such as dropout) require a JAX PRNG key to generate random masks during training. Because JAX is purely functional, key propagation must be handled explicitly:

• Training Mode: To apply stochastic behavior (e.g., dropout), pass a jax.random.PRNGKey to the model call via the key argument, and set inference=False.
• Inference Mode: To compute the deterministic forward pass, pass inference=True (or omit the key argument). When inference=True, the model operates in evaluation mode and disables all dropout masking.

Example: Running FNO with Dropout

import jax.numpy as jnp
import jax.random as jr
from neojax.models import FNO

key = jr.key(0)
model_key, train_key = jr.split(key)

# Initialize FNO with channel MLP dropout
model = FNO(
    key=model_key,
    in_channels=1,
    out_channels=1,
    hidden_channels=32,
    n_layers=4,
    modes=(16,),
    channel_mlp_dropout=0.1,
)

x = jnp.ones((1, 64))

# Training forward pass (stochastic, requires key)
out_train = model(x, key=train_key, inference=False)

# Inference forward pass (deterministic, no key needed)
out_eval = model(x, inference=True)