RNO

Recurrent Neural Operator (RNO)

neojax.models.rno.RNO

General Recurrent Neural Operator architecture.

Essentially wraps another Neural Operator and applies it recursively to its outputs for \(t\) timesteps over a time interval \([0, \dots, T]\).

The recurrent step is model-agnostic. It handles three kinds of input:

• the state \(\hat{u}_n\), which is the lax.scan carry,
• the per-step forcing \(f_{n+1}\), scanned over its leading time axis,
• static inputs (key, inference and any extra keyword arguments of the wrapped model), constant across steps and not scanned.

How state and forcing are turned into the wrapped model's input is handled by combine_fn. The default concatenates them along the channel axis, which fits FNO, TFNO, UNO and any other model taking a single multichannel array. The wrapped model then needs in_channels = state_channels + forcing_channels and out_channels = state_channels.

Parameters:

Name	Type	Description	Default
model	BaseNO	A suitable model instance to wrap.	required
time_interval	tuple[float | int, float | int]	Time interval [start, end] of trajetory.	required
n_timesteps	int	Number of timesteps in the time trajectory.	required
combine_fn	Callable[[Inexact[Array, 'c ...'], PyTree[Inexact[Array, 'f ...']] | None], Inexact[Array, ...]]	Function that determines how model inputs and additional optional per-step inputs are combined. For most models (e.g. FNO-style, MgNO), this should be a concatenation along the channel dimension: lambda x, forcing: jnp.concat((x, forcing), axis=0). This is also the default.	concat_channels
checkpoint_timesteps	bool	Whether to recompute linearizations at each autoregressive timestep. Can decrease memory at the cost of increased computation time. Default is False.	False

Raises:

Type	Description
ValueError	If n_timesteps is smaller 1.
ValueError	If time_interval is invalid.

Internal Attributes

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

• model (BaseNO): Suitable model instance to wrap.
• n_timesteps (int): Number of timesteps in the time trajectory.
• dt (float): Per step time delta.
• combine_fn (Callable[ [Inexact[Array, "c ..."], PyTree[Inexact[Array, "_ ..."]] | None], Inexact[Array, "..."], ]): Function that determines how model inputs and additional optional per-step inputs are combined.
• checkpoint_timesteps (bool): Whether to recompute linearizations at each autoregressive timestep. Default is False.

Cite

Recurrent Neural Operators: Stable Long-Term PDE Prediction

@article{ye2025recurrent,
    title={Recurrent neural operators: stable long-term PDE prediction},
    author={Ye, Zaijun and Zhang, Chen-Song and Wang, Wansheng},
    journal={arXiv preprint arXiv:2505.20721},
    year={2025}
}

Time and Memory complexity

Computation time and memory consumption increase with the number of timesteps \(t\). Even with checkpoint_timesteps, expect time per epoch and memory usage to be far greater than for a non-recurrent model.

step(x: Inexact[Array, 'c ...'], f_t: PyTree[Inexact[Array, '_ ...']] | None = None, *, key: PRNGKeyArray | None = None, inference: bool = False, **static: typing.Any) -> Inexact[Array, 'c ...']

Computes a single timestep.

Given initial \(u_0\) and optional per-step forcing \(f_{n+1}\), this computes the explicit Euler update

\[ \hat{u}_{n+1} = \hat{u}_{n} + \Delta t \, \mathcal{G}_{\theta}(\hat{u}_{n}, f_{n+1}), \; n = 0, 1, \dots, N - 1. \]

Parameters:

Name	Type	Description	Default
x	Inexact[Array, 'c ...']	Current state of shape (in_c, d1, ..., dN), where N are the number of spatial dimensions.	required
f_t	PyTree[Inexact[Array, '_ ...']] | None	Pytree of this step's forcing, passed to combine_fn together with x. Default is None.	None
key	PRNGKeyArray | None	PRNG key used for the wrapped model. Default is None.	None
inference	bool	If True, training-only logic of the wrapped model is disabled. Default is False.	False
static	Any	Static extra keyword arguments passed unchanged to the wrapped model (e.g. GeoFNO coordinates).	required

Returns:

Type	Description
Inexact[Array, 'c ...']	The state \(\hat{u}_{n+1}\) of the same shape as x.

Raises:

Type	Description
ValueError	If model output's shape of step n doesn't match the model input's shape of step n + 1.

__call__(x: Inexact[Array, 'c ...'], forcing: PyTree[Inexact[Array, 't _ ...']] | None = None, *, key: PRNGKeyArray | None = None, inference: bool = False, **static: typing.Any) -> Inexact[Array, 't c ...']

Forward pass of the RNO.

Applies step recursively with given initial \(\hat{u}_0\) and optional per-step parameters. The outputs across all timesteps are stacked into a single output with leading time axis.

Parameters:

Name	Type	Description	Default
x	Inexact[Array, 'c ...']	Initial state of shape (c, d1, ..., dN), where N are the number of spatial dimensions.	required
forcing	PyTree[Inexact[Array, 't _ ...']] | None	Pytree of per-step forcing arrays, each of shape (n_timesteps, f_c, d1, ..., dN). Slice t is used at step t. Time-constant fields that should be concatenated to the input need to be broadcast along the time axis, e.g. with neojax.utils.tree_add_broadcast_to. Default is None, no forcing.	None
key	PRNGKeyArray | None	PRNG key used for the wrapped model. Split into one key per timestep. Default is None.	None
inference	bool	If True, training-only logic of the wrapped model is disabled. Default is False.	False
static	Any	Static extra keyword arguments passed unchanged to the wrapped model at every step (e.g. GeoFNO coordinates). They are not scanned over.	required

Returns:

Type	Description
Inexact[Array, 't c ...']	The predicted field of shape (n_timesteps, out_channels, d1, ..., dN).

Raises:

Type	Description
ValueError	If any leaf of the forcing tree doesn't have a leading time axis of size n_timesteps.