neuralop.models.RNO

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

N-Dimensional Recurrent Neural Operator.

The RNO has an identical architecture to the finite-dimensional GRU, with the exception that linear matrix-vector multiplications are replaced by Fourier layers (Li et al., 2021), and for regression problems, the output nonlinearity is replaced by a SELU activation.

The operation of the GRU is as follows: z_t = sigmoid(W_z x + U_z h_{t-1} + b_z) r_t = sigmoid(W_r x + U_r h_{t-1} + b_r) hat h_t = selu(W_h x_t + U_h (r_t * h_{t-1}) + b_h) h_t = (1 - z_t) * h_{t-1} + z_t * hat h_t,

where * is element-wise, the b_i’s are bias functions, and W_i, U_i are linear Fourier layers.

Parameters:

n_modesTuple[int, …]

Number of modes to keep in Fourier Layer, along each dimension. The dimensionality of the RNO is inferred from len(n_modes). n_modes must be larger enough but smaller than max_resolution//2 (Nyquist frequency)

in_channelsint

Number of input channels in input function. Determined by the problem.

out_channelsint

Number of output channels in output function. Determined by the problem.

hidden_channelsint

Width of the RNO (i.e. number of channels). This significantly affects the number of parameters of the RNO. For 1D problems, 24-48 channels are typically sufficient. For 2D/3D problems without Tucker/CP/TT factorization, start with 16-32 channels to avoid excessive parameters. Increase if more expressivity is needed. Update lifting_channel_ratio and projection_channel_ratio accordingly since they are proportional to hidden_channels.

n_layersint, optional

Number of RNO layers to use. Default: 4

rno_skipbool, optional

Whether to use skip connections between RNO layers. When True, adds the input to each layer’s output: x_{l+1} = x_l + RNOBlock(x_l). Default: False. Unlike FNO where skip connections improve gradient flow, RNO’s recurrent structure already provides temporal gradient pathways, so skip connections are optional and may not always improve performance.

Methods

forward(x[, init_hidden_states, ...])	Forward pass for the Recurrent Neural Operator.
predict(x, num_steps[, grid_function])	Autoregressively predict future time steps.

Other Parameters:

lifting_channel_ratioNumber, optional

Ratio of lifting channels to hidden_channels. The number of lifting channels in the lifting block of the RNO is lifting_channel_ratio * hidden_channels (e.g. default 2 * hidden_channels).

projection_channel_ratioNumber, optional

Ratio of projection channels to hidden_channels. The number of projection channels in the projection block of the RNO is projection_channel_ratio * hidden_channels (e.g. default 2 * hidden_channels).

positional_embeddingUnion[str, nn.Module], optional

Type of positional embedding to use. Options: - “grid”: Use default grid-based positional embedding - GridEmbeddingND: Custom N-dimensional grid embedding - GridEmbedding2D: Custom 2D grid embedding (only for 2D problems) - None: No positional embedding Default: None

non_linearitynn.Module, optional

Non-Linear activation function module to use. Default: F.gelu

normLiteral[“ada_in”, “group_norm”, “instance_norm”], optional

Normalization layer to use. Options: “ada_in”, “group_norm”, “instance_norm”, None. Default: None

complex_databool, optional

Whether the data is complex-valued. If True, initializes complex-valued modules. Default: False

use_channel_mlpbool, optional

Whether to use an MLP layer after each FNO block. Default: True

channel_mlp_dropoutfloat, optional

Dropout parameter for ChannelMLP in FNO Block. Default: 0

channel_mlp_expansionfloat, optional

Expansion parameter for ChannelMLP in FNO Block. Default: 0.5

channel_mlp_skipLiteral[“linear”, “identity”, “soft-gating”, None], optional

Type of skip connection to use in channel-mixing mlp. Options: “linear”, “identity”, “soft-gating”, None. Default: “soft-gating”

fno_skipLiteral[“linear”, “identity”, “soft-gating”, None], optional

Type of skip connection to use in FNO layers. Options: “linear”, “identity”, “soft-gating”, None. Default: “linear”

return_sequencesbool, optional

Whether the final RNO layer returns the full sequence of hidden states or just the final state. Intermediate layers always return sequences. Default: False

resolution_scaling_factorUnion[Number, List[Number]], optional

Layer-wise factor by which to scale the domain resolution of function. Options: - None: No scaling - Single number n: Scales resolution by n at each layer - List of numbers [n_0, n_1,…]: Scales layer i’s resolution by n_i Default: None

domain_paddingUnion[Number, List[Number]], optional

Percentage of padding to use. Options: - None: No padding - Single number: Percentage of padding to use along all dimensions - List of numbers [p1, p2, …, pN]: Percentage of padding along each dimension Default: None

fno_block_precisionstr, optional

Precision mode in which to perform spectral convolution. Options: “full”, “half”, “mixed”. Default: “full”.

stabilizerstr, optional

Whether to use a stabilizer in FNO block. Options: “tanh”, None. Default: None. stabilizer greatly improves performance in the case fno_block_precision=’mixed’.

max_n_modesTuple[int, …], optional

Maximum number of modes to use in Fourier domain during training. None means that all the n_modes are used. Tuple of integers: Incrementally increase the number of modes during training. This can be updated dynamically during training.

factorizationstr, optional

Tensor factorization of the FNO layer weights to use. Options: “None”, “Tucker”, “CP”, “TT” Other factorization methods supported by tltorch. Default: None

rankfloat, optional

Tensor rank to use in factorization. Default: 1.0 Set to float <1.0 when using TFNO (i.e. when factorization is not None). A TFNO with rank 0.1 has roughly 10% of the parameters of a dense FNO.

fixed_rank_modesbool, optional

Whether to not factorize certain modes. Default: False

implementationstr, optional

Implementation method for factorized tensors. Options: “factorized”, “reconstructed”. Default: “factorized”

decomposition_kwargsdict, optional

Extra kwargs for tensor decomposition (see tltorch.FactorizedTensor). Default: {}

separablebool, optional

Whether to use a separable spectral convolution. Default: False

preactivationbool, optional

Whether to compute FNO forward pass with resnet-style preactivation. Default: False

conv_modulenn.Module, optional

Module to use for FNOBlock’s convolutions. Default: SpectralConv

enforce_hermitian_symmetrybool, optional

Whether to enforce Hermitian symmetry conditions when performing inverse FFT for real-valued data. Only used in SpectralConv; ignored otherwise. When True, explicitly enforces that the 0th frequency and Nyquist frequency are real-valued before calling irfft. When False, relies on cuFFT’s irfftn to handle symmetry automatically, which may fail on certain GPUs or input sizes, causing line artifacts. By default True.

References

[RNO]

Liu-Schiaffini, M., Singer, C. E., Kovachki, N., Schneider, T., Azizzadenesheli, K., & Anandkumar, A. (2023). Tipping point forecasting in non-stationary dynamics on function spaces. arXiv preprint arXiv:2308.08794.

forward(x, init_hidden_states=None, return_hidden_states=False, keep_states_padded=False)[source]

Forward pass for the Recurrent Neural Operator.

Parameters:

xtorch.Tensor

Input tensor with shape (batch, timesteps, in_channels, *spatial_dims), where len(spatial_dims) == self.n_dim. The channel dimension MUST be at index 2 and the time dimension MUST be at index 1.

init_hidden_stateslist[torch.Tensor] | None

Optional list of per-layer initial hidden states. Each tensor should have shape (batch, hidden_channels, *spatial_dims_h), where spatial_dims_h are the unpadded spatial dimensions at each layer (accounting for resolution scaling). If domain_padding is enabled, hidden states will be automatically padded like the input x, unless keep_states_padded is True. If None, all hidden states are initialized internally.

return_hidden_statesbool, optional

Whether to return the final hidden states for each layer in addition to the prediction. Default: False

keep_states_paddedbool, optional

If True, treats provided init_hidden_states as already padded (no additional padding applied) and returns final_hidden_states in padded form (no unpadding). Use this for efficient autoregressive loops to preserve boundary information. Default: False

Returns:

predtorch.Tensor

Output tensor with shape (batch, out_channels, *spatial_dims_out).

final_hidden_stateslist[torch.Tensor], optional

Returned only if return_hidden_states=True.

predict(x, num_steps, grid_function=None)[source]

Autoregressively predict future time steps.

Performs autoregressive rollout by iteratively feeding predictions back as input. At each step, the model predicts the next time step from the current input, then uses that prediction as input for the subsequent prediction.

Parameters:

xtorch.Tensor

Initial input sequence with shape (batch, timesteps, in_channels, *spatial_dims). The last time step of this sequence serves as the starting point for predictions.

num_stepsint

Number of future time steps to predict autoregressively.

grid_functioncallable, optional

Function that generates positional embeddings (e.g., spatial coordinates) to concatenate with predictions before feeding them back as input. Should have signature grid_function(shape, device) -> torch.Tensor, where the returned tensor has shape matching the input requirements. Use this when your model was trained with concatenated positional information. Default: None

Returns:

torch.Tensor

Predicted sequence with shape (batch, num_steps, out_channels, *spatial_dims), containing the autoregressively generated future states.

Notes

This method maintains hidden states across prediction steps, enabling the RNO to leverage its recurrent structure during autoregressive generation.