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.

Paper: .. [R943463ff80d7-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.

Methods

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

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.