Geometry Informed Neural Operator (GiNO)

neojax.models.gino.GiNO

General n-dimensional Geometry-Informed Neural Operator (GiNO).

GiNO learns a map between functions given on arbitrary point clouds or meshes. It combines the flexibility of the Graph Neural Operator (GNO) with the efficiency of the FNO. The forward pass is

1. an input GNOBlock that integrates the input function over the neighborhoods of the points of a regular latent grid,
2. optional concatenation of additional latent grid features,
3. a lifting PointwiseMLP and the FNOBlocks on the latent grid,
4. an output GNOBlock that integrates the latent function over the neighborhoods of the output query points,
5. a final projection PointwiseMLP to the output channels.

The output GNO can be turned into a mollified GNO by selecting a weighting function (gno_weighting_function). The output is then differentiable in the query positions x_out.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key for parameter initialization.	required
in_channels	int	Number of channels of the input function.	required
out_channels	int	Number of channels of the output function.	required
modes	Sequence[int]	Number of Fourier modes to retain across each dimension of the latent grid, e.g. (16, 16, 16). Its length is the dimension of the latent grid. Must be smaller than half of the latent grid resolution (Nyquist frequency).	required
hidden_channels	int	Hidden channel dimension of the latent FNO.	64
n_layers	int	Number of consecutive FNO blocks.	4
coord_dim	int	Dimension of the physical coordinates of the input points and output queries. Default is 3.	3
latent_feature_channels	int | None	Number of channels of an optional latent feature map, concatenated to the output of the input GNO before the lifting. Default is None, no latent features.	None
in_gno_radius	float	Radius of the neighborhoods of the input GNO. Larger radii mean more interaction at a higher cost. Default is 0.033.	0.033
out_gno_radius	float	Radius of the neighborhoods of the output GNO. Default is 0.033.	0.033
in_gno_max_neighbors	int	Maximum number of neighbors per latent point in the input GNO. If more points are inside the radius, only the nearest are used. Default is 32.	32
out_gno_max_neighbors	int	Maximum number of neighbors per output query in the output GNO. Default is 32.	32
in_gno_transform_type	Literal['linear_kernelonly', 'linear', 'nonlinear_kernelonly', 'nonlinear']	Integral transform of the input GNO, see IntegralTransform. Default is "linear".	'linear'
out_gno_transform_type	Literal['linear_kernelonly', 'linear', 'nonlinear_kernelonly', 'nonlinear']	Integral transform of the output GNO. Default is "linear".	'linear'
in_gno_out_channels	int | None	Number of output channels of the input GNO, i.e. input channels of the latent FNO (without latent features). Must equal in_channels for the "linear" and "nonlinear" transforms. Default is None, which is in_channels.	None
gno_weighting_function	Optional[Literal['bump', 'half_cos', 'quadr', 'quartic', 'octic']]	Weighting function of the output GNO for a mollified GNO, see make_weighting_fn. Default is None.	None
gno_weight_function_scale	float	Factor by which to scale the weights of gno_weighting_function. Default is 1.0.	1.0
in_gno_pos_embedding_type	Optional[Literal['transformer', 'nerf']]	Type of the sinusoidal embedding of the points in the input GNO, "transformer", "nerf" or None for no embedding. Default is "transformer".	'transformer'
out_gno_pos_embedding_type	Optional[Literal['transformer', 'nerf']]	Type of the sinusoidal embedding of the points in the output GNO. Default is "transformer".	'transformer'
gno_embed_channels	int	Number of frequencies of the sinusoidal embeddings per coordinate. Default is 32.	32
gno_embed_max_positions	int	max_positions of the "transformer" embeddings. Default is 10000.	10000
in_gno_channel_mlp_hidden_layers	Sequence[int]	Widths of the hidden layers of the kernel network of the input GNO. Default is (80, 80, 80).	(80, 80, 80)
out_gno_channel_mlp_hidden_layers	Sequence[int]	Widths of the hidden layers of the kernel network of the output GNO. Default is (512, 256).	(512, 256)
gno_activation	Callable	Activation function of the GNO kernel networks. Default is jax.nn.gelu.	gelu
out_gno_tanh	bool	Whether to apply tanh to the latent function before the output GNO, to stabilize it. Default is False.	False
gno_query_chunk_size	int | None	Number of query points the GNO neighbor searches and integral transforms process at a time to limit peak memory. Default is None, all points at once.	None
activation	Callable	Activation function of the lifting, FNO blocks and projection. Default is jax.nn.gelu.	gelu
use_channel_mlp	bool	Whether to apply a pointwise channel MLP after each FNO block. Default is True.	True
local_operator	Optional[Literal['linear', 'soft-gating', 'identity']]	Type of local operator inside FNO blocks. Can be "linear", "soft-gating", "identity", or None. Default is "linear".	'linear'
use_local_operator_bias	bool	Whether to use a bias term in the FNO blocks local operator. Default is False.	False
channel_mlp_residual	Optional[Literal['linear', 'soft-gating', 'identity']]	Type of skip connection around channel MLPs. Can be "linear", "soft-gating", "identity", or None. Default is "soft-gating".	'soft-gating'
channel_mlp_expansion	float | None	Expansion factor for hidden dimension in the channel MLPs. Default is 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. Default is 0.0.	0.0
normalization	Optional[Literal['layer', 'instance', 'group']]	Type of normalization to use in FNO blocks. Can be "layer", "instance", "group" or None. Default is None.	None
norm_groups	int	Number of groups for group normalization. Default is 1.	1
use_fno_residual	bool	Whether to use a Resnet-style residual connection around FNO blocks. Default is False.	False
preactivation	bool	Whether to use pre-activation style blocks. Default is False.	False
n_lift_layers	int	Number of layers in the lifting MLP. Default is 2.	2
n_proj_layers	int	Number of layers in the projection MLP. Default is 2.	2
lift_channel_ratio	float	Ratio of lifting channels to hidden_channels. Default is 2.0.	2.0
proj_channel_ratio	float	Ratio of projection channels to hidden_channels. Default is 4.0.	4.0
enforce_hermitian_symmetry	bool	Whether to enforce hermitian symmetry on the outputs of the spectral convolutions before calling irfftn. Default is True.	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 latent data is complex valued. If True, uses full FFT. 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, by default 'auto'. If 'auto', uses (2 / (in_channels * out_channels)) ** 0.5.	'auto'
factorization	Union[BaseTensor, Literal['tucker', 'cp', 'tt'], NoneType]	Tensor factorization type of the spectral weights. Can be a BaseTensor instance, a string ("tucker", "cp", "tt"), or None. Default is None.	None
implementation	Literal['reconstructed', 'factorized']	Weight reconstruction mode. Can be "reconstructed" or "factorized". Default is "factorized".	'factorized'
separable	bool	Whether to use separable implementation of contraction. Default is False.	False

Raises:

Type	Description
ValueError	If modes is empty or contains non-positive entries.
ValueError	If gno_weighting_function is invalid.
ValueError	If the channels of the input GNO don't match its transform type, see GNOBlock.

Internal Attributes

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

• in_gno (GNOBlock): GNO from the input points to the latent grid.
• lifting (PointwiseMLP): The PointwiseMLP lifting the latent function to hidden_channels.
• fno_blocks (FNOBlocks): The FNOBlocks on the latent grid.
• out_gno (GNOBlock): GNO from the latent grid to the output queries.
• projection (PointwiseMLP): The PointwiseMLP projecting to out_channels.
• latent_feature_channels (int): Number of latent feature channels.
• latent_dim (int): Dimension of the latent grid.
• out_gno_tanh (bool): Whether tanh is applied before the output GNO.

Cite

Geometry-Informed Neural Operator for Large-Scale 3D PDEs

@inproceedings{li2023geometry,
    title={Geometry-Informed Neural Operator for Large-Scale 3D PDEs},
    author={Li, Zongyi and Kovachki, Nikola and Choy, Chris and Li, Boyi
        and Kossaifi, Jean and Otta, Shourya and Nabian, Mohammad Amin
        and Stadler, Maximilian and Hundt, Christian and
        Azizzadenesheli, Kamyar and Anandkumar, Anima},
    booktitle={Advances in Neural Information Processing Systems},
    year={2023}
}

Neural Operator: Graph Kernel Network for Partial Differential Equations

@article{li2020neural,
    title={Neural operator: Graph kernel network for
        partial differential equations},
    author={Li, Zongyi and Kovachki, Nikola and Azizzadenesheli,
        Kamyar and Liu, Burigede and Bhattacharya, Kaushik
        and Stuart, Andrew and Anandkumar, Anima},
    journal={arXiv preprint arXiv:2003.03485},
    year={2020}
}

Enabling Automatic Differentiation with Mollified Graph Neural Operators

@article{lin2025enabling,
    title={Enabling automatic differentiation with mollified graph neural operators},
    author={Lin, Ryan Y and Berner, Julius and Duruisseaux,
    Valentin and Pitt, David and Leibovici, Daniel and Kossaifi,
    Jean and Azizzadenesheli, Kamyar and Anandkumar, Anima},
    journal={arXiv preprint arXiv:2504.08277},
    year={2025}
}

vmap and fixed geometries

To batch the vmap call, map over the function values u (and latent_features) and share the geometry with in_axes=None

out = jax.vmap(
    model, in_axes=(0, None, None, None)
)(u_batch, x_in, x_latent, x_out)

The geometries can also be batched, since the neighbor search is jit- and vmap-compatible. For a fixed geometry, precompute the neighbors once (radius_neighbor_search) and pass them with in_neighbors and out_neighbors to avoid searching in every call.

Neighbor truncation

The GNOs consider at most in_gno_max_neighbors and out_gno_max_neighbors neighbors per point (the nearest ones). Choose these large enough for your point densities and radii. With a gno_weighting_function, truncated neighbors can have nonzero weights, which breaks the smoothness of the output. Then choose the values large enough that truncation does not occur.

Extent of the latent grid

The mollified GNO paper pads the latent domain, so that the latent function is supported on the geometry extended by the radius, \(D + B_r(0)\). Let x_latent extend about in_gno_radius and out_gno_radius beyond the geometry, such that the GNOs see a full neighborhood at the boundary of the domain.

__call__(u: Float[Array, 'in_c n'] | None, *, x_in: Float[Array, 'n d'], x_latent: Float[Array, '... d'], x_out: Float[Array, 'm d'] | dict[str, Float[Array, '_ d']] | None = None, latent_features: Float[Array, 'lat_c ...'] | None = None, in_quadrature_weights: Float[Array, n] | None = None, latent_quadrature_weights: Float[Array, ...] | None = None, in_neighbors: neojax.nn.neighbor_search.Neighbors | None = None, out_neighbors: neojax.nn.neighbor_search.Neighbors | dict[str, neojax.nn.neighbor_search.Neighbors] | None = None, key: PRNGKeyArray | None = None, inference: bool = False) -> Float[Array, 'out_c m'] | dict[str, Float[Array, 'out_c _']]

Forward pass of the GINO model.

Parameters:

Name	Type	Description	Default
u	Float[Array, 'in_c n'] | None	Input function sampled at x_in, shaped (in_channels, n). Can be None if the input GNO transform doesn't need it ("linear", "linear_kernelonly"), in which case the model only sees the geometry.	required
x_in	Float[Array, 'n d']	Coordinates of the n input points, shaped (n, coord_dim).	required
x_latent	Float[Array, '... d']	Coordinates of the regular latent grid, shaped (g1, ..., gk, coord_dim), where k = len(modes). The grid is typically a uniform grid on the unit cube or on the bounding box of the geometry. The physical radii of the GNOs are in the units of these coordinates.	required
x_out	Float[Array, 'm d'] | dict[str, Float[Array, '_ d']] | None	Coordinates of the output queries, shaped (m, coord_dim). Pass a dict of such arrays to query at several sets of points, e.g. the volume and the surface of the domain, which returns a dict with the same keys. Default is None, which queries at x_in.	None
latent_features	Float[Array, 'lat_c ...'] | None	Optional features on the latent grid, shaped (latent_feature_channels, g1, ..., gk). Required iff latent_feature_channels was set.	None
in_quadrature_weights	Float[Array, n] | None	Optional quadrature weights of the n input points x_in for the input GNO, see IntegralTransform. Default is None.	None
latent_quadrature_weights	Float[Array, ...] | None	Optional quadrature weights of the latent grid points for the output GNO, shaped (g1, ..., gk) like the latent grid, e.g. the cell volume. Default is None.	None
in_neighbors	Neighbors | None	Optional precomputed neighbors in x_in of the flattened latent points x_latent. Default is None, the neighbors are searched in every call.	None
out_neighbors	Neighbors | dict[str, Neighbors] | None	Optional precomputed neighbors in the flattened x_latent of the output queries. A dict with the keys of x_out if x_out is a dict. Default is None.	None
key	PRNGKeyArray | None	PRNG key used for dropout masks in FNO blocks.	None
inference	bool	If True, dropout is disabled.	False

Returns:

Type	Description
Float[Array, 'out_c m'] | dict[str, Float[Array, 'out_c _']]	The output function sampled at x_out, shaped (out_channels, m), or a dict of such arrays if x_out is a dict.

Raises:

Type	Description
ValueError	If x_latent doesn't have len(modes) grid axes.
ValueError	If latent_features are missing, unexpected or have a shape other than the latent grid.
ValueError	If out_neighbors is a dict but x_out is not or the keys of both differ.
ValueError	If latent_quadrature_weights doesn't have the shape of the latent grid.