UNO

U-shaped FNO (UNO)

neojax.models.uno.UNO

U-shaped Neural Operator.

The architecture is described in the reference publication. It extends the standard FNO with a U-shaped encoder-decoder structure and horizontal skip connections between corresponding encoder and decoder layers.

Each FNO layer has its own number of output channels, Fourier modes, and resolution scaling factor. Horizontal skip connections concatenate features from an encoder layer to a decoder layer along the channel dimension. An optional skip-connection projection is applied before concatenation.

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	Initial width of the UNO. This significantly affects the number of parameters of the UNO. Good starting point can be 64, and then increased if more expressivity is needed. (Update lift_channel_ratio and proj_channel_ratio accordingly).	required
uno_out_channels	Sequence[int]	Number of output channels of each Fourier layer, e.g. for a 5 layer UNO uno_out_channels can be [32, 64, 64, 64, 32].	required
uno_modes	Sequence[Sequence[int]]	Number of Fourier modes to use in integral operation of each Fourier layer (along each dimension). For example in a 5 layer UNO with 2D inputs: [[5, 5], [5, 5], [5, 5], [5, 5], [5, 5]].	required
uno_scalings	Sequence[float | int] | Sequence[Sequence[float | int]]	Scaling factor for each Fourier layer. For e.g. a 5 layer UNO with 2D inputs, the uno_scalings can be [[0.5, 0.5], [0.5, 0.5], [1, 1], [2, 2], [2, 2]]. If inner elements are scalars, the same scaling is applied to all spatial dims.	required
n_fno_layers	int	Number of Fourier layers. Default is 4.	4
n_lift_layers	int	Number of lifting layers. Default is 2.	2
n_proj_layers	int	Number of projection layers. Default is 2.	2
lift_channel_ratio	float	Ratio of lifting channels to hidden_channels. The number of lifting channels in the lifting block of the UNO 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 UNO is projection_channel_ratio * hidden_channels (e.g. default 2 * hidden_channels).	2.0
positional_embedding	Union[GridEmbeddingNd, Literal['grid'], NoneType]	Positional embedding to apply to last channels of raw input before passing through UNO. Default is None. Pass "grid" for a regular GridEmbeddingNd.	None
horizontal_skip_map	dict[int, int] | None	A dictionary {b: a, ...} denoting horizontal skip connections from the a-th layer to the b-th layer. If None, default symmetric skip connections are applied (mirroring encoder layers to decoder layers). Layer indices are zero-based.	None
use_channel_mlp	bool	Whether to apply a pointwise channel MLP after each FNO block. Defaults to True.	True
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
channel_mlp_dropout	float	Dropout probability applied after each layer (except the last) of the channel MLPs. Defaults to 0.0.	0.0
channel_mlp_residual	Optional[Literal['linear', 'soft-gating', 'identity']]	Type of skip connection around channel MLPs. Can be "linear", "soft-gating", "identity", or None. Defaults to "soft-gating".	'soft-gating'
use_fno_residual	bool	Whether to use residual connection around FNO blocks. Default is True.	True
activation	Callable	Activation function used within the layers. Defaults to jax.nn.gelu.	gelu
normalization	Optional[Literal['layer', 'instance', 'group']]	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
preactivation	bool	Whether to use pre-activation style blocks. Default is False.	False
local_operator	Optional[Literal['linear', 'soft-gating', 'identity']]	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
horizontal_residual	Optional[Literal['linear', 'soft-gating', 'identity']]	Type of skip connection to use for horizontal skip paths. Can be "linear", "soft-gating", "identity", or None. Defaults to "linear".	'linear'
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
enforce_hermitian_symmetry	bool	Whether to enforce hermitian symmetry on the outputs of the spectral convolutions before calling irfftn. Default is True. If set to False, cuFFT on GPU may cause line artifacts when calling irfftn.	True
fft_norm	Optional[Literal['forward', 'backward', 'ortho']]	FFT normalization. Can be None, "backward", "ortho" or "forward". Default is "forward".	'forward'
is_complex_data	bool	Whether the input data is complex valued. Default is False.	False
ranks	int | Sequence[int] | None	Number of ranks to contract the spectral tensors to. Default is None.	None
init_std	Union[float, Literal['auto']]	Standard deviation to use for weight initialization. Default is "auto".	'auto'
factorization	Union[BaseTensor, Literal['tucker', 'cp', 'tt'], NoneType]	Tensor factorization type. Default is None.	None
implementation	Literal['reconstructed', 'factorized']	Weight reconstruction mode. Default is "factorized".	'factorized'
separable	bool	Whether to use separable implementation of contraction. Default is False.	False

Internal Attributes

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

• positional_embedding (GridEmbeddingNd | None): Positional embedding.
• lifting (PointwiseMLP): Maps inputs to hidden dimension.
• fno_blocks (tuple[FNOBlock, ...]): Per-layer FNO blocks with layer-specific channels, modes, and scalings.
• channel_mlps (tuple[PointwiseMLP, ...] | None): Per-layer channel MLPs.
• channel_mlp_residuals (tuple | None): Per-layer channel MLP residual connections.
• horizontal_skip_connections (dict[int, Flattened1dConv | SoftGating | eqx.nn.Identity | None]): Skip connections for horizontal paths.
• projection (PointwiseMLP): Maps latent features to output channels.
• padding (DomainPadding | None): Percentage of domain padding to use.
• n_fno_layers (int): Number of FNO layers.
• end_to_end_scaling_factor (list[int | float]): Cumulative scaling factor per spatial dim.
• horizontal_skip_map (dict[int, int] | None): Maps decoder layer indices to encoder layer indices for skip connections.

Examples:

import jax.numpy as jnp
import jax.random as jr
from neojax.models.uno import UNO

key = jr.key(0)
model = UNO(
    key=key,
    in_channels=3,
    out_channels=1,
    hidden_channels=16,
    uno_out_channels=[16, 32, 32, 16],
    uno_modes=[[8, 8], [8, 8], [8, 8], [8, 8]],
    uno_scalings=[[0.5, 0.5], [0.5, 0.5], [2, 2], [2, 2]],
    n_fno_layers=4,
)
x = jnp.ones((3, 32, 32))
out = model(x)

Cite

U-no: U-shaped neural operators

@article{rahman2022u,
    title={U-no: U-shaped neural operators},
    author={Rahman, Md Ashiqur and Ross,
    Zachary E and Azizzadenesheli, Kamyar},
    journal={arXiv preprint arXiv:2204.11127},
    year={2022}
}

__call__(x: Inexact[Array, 'in_c ...'], *, key: PRNGKeyArray | None = None, inference: bool = False) -> Inexact[Array, 'out_c ...']

Forward pass of the UNO model.

The input is lifted, passed through the U-shaped FNO layers with horizontal skip connections, and projected to the output.

Parameters:

Name	Type	Description	Default
x	Inexact[Array, 'in_c ...']	Input array of shape (in_channels, d1, ..., dN).	required
key	PRNGKeyArray | None	PRNG key used for dropout masks in channel MLPs.	None
inference	bool	If True, dropout is disabled.	False

Returns:

Type	Description
Inexact[Array, 'out_c ...']	The predicted field of shape (out_channels, d1', ..., dN'), where d_i' depends on the cumulative scaling factors.