GNO Block

GNO Block

neojax.nn.gno_block.GNOBlock

Graph Neural Operator block.

A resolution-invariant operator that maps a function sampled on one set of points \(y\) to a function sampled on another set of points \(x\). The output at \(x\) is a kernel integral over those input points that are within a ball of fixed radius around \(x\), see IntegralTransform for the available transforms.

The points can be optionally embedded with a SinusoidalEmbedding before they enter the kernel network. The neighbors are always searched with the raw coordinates.

With a weighting_fn the block is a mollified GNO. The hard ball is replaced by a smooth weight, so the output is differentiable in the query positions x. The weights are computed from the raw coordinates inside the block, so the derivatives w.r.t. x are also correct for precomputed neighbors.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key for parameter initialization.	required
in_channels	int	Number of channels of the input function. Used by the nonlinear transforms (as input of the kernel) and to validate the "linear" transforms.	required
out_channels	int	Number of channels of the output function.	required
coord_dim	int	Dimension of the domain on which \(x\) and \(y\) are defined.	required
radius	float	Radius of the neighborhoods.	required
max_neighbors	int	Maximum number of neighbors per output point. Neighbors beyond the nearest max_neighbors are dropped. Default is 32.	32
transform_type	Literal['linear_kernelonly', 'linear', 'nonlinear_kernelonly', 'nonlinear']	Integral transform to compute. Default is "linear".	'linear'
reduction	Literal['sum', 'mean']	How to aggregate over a neighborhood, by summing ("sum") or averaging ("mean"). Default is "sum".	'sum'
weighting_fn	Callable[[Float[Array, ...]], Float[Array, ...]] | None	Optional mollifier applied to the squared distances between the points and their neighbors, see make_weighting_fn. Forces the "sum" reduction. Default is None.	None
pos_embedding_type	Optional[Literal['transformer', 'nerf']]	Type of the SinusoidalEmbedding of the points, "transformer" or "nerf". None disables the embedding. Default is "transformer".	'transformer'
pos_embedding_channels	int	Number of frequencies of the embedding per coordinate. Default is 32.	32
pos_embedding_max_positions	int	max_positions of the "transformer" embedding. Default is 10000.	10000
channel_mlp_hidden_layers	Sequence[int]	Widths of the hidden layers of the kernel network. The input and output widths are inferred. Default is (128, 256, 128). Not used if channel_mlp is passed.	(128, 256, 128)
activation	Callable	Activation function of the kernel network. Default is jax.nn.gelu.	gelu
channel_mlp	PointwiseMLP | None	Optional custom kernel network. Must only mix channels to keep the integral discretization invariant and have matching input and output channels (see the kernel_in_channels attribute). If passed, channel_mlp_hidden_layers and activation are ignored.	None
query_chunk_size	int | None	Number of output points to process at a time in the neighbor search and the transform to limit peak memory. Default is None, all points at once.	None

Raises:

Type	Description
ValueError	If in_channels, out_channels, coord_dim, max_neighbors or pos_embedding_channels are not positive integers.
ValueError	If radius is not positive.
ValueError	If transform_type is "linear" or "nonlinear" and in_channels differs from out_channels.
ValueError	If channel_mlp has the wrong number of input or output channels.

Internal Attributes

• pos_embedding (SinusoidalEmbedding | None): Embedding of the points.
• integral_transform (IntegralTransform): The kernel integral transform.
• in_channels (int): Number of channels of the input function.
• out_channels (int): Number of channels of the output function.
• coord_dim (int): Dimension of the domain.
• radius (float): Radius of the neighborhoods.
• max_neighbors (int): Maximum number of neighbors per output point.
• query_chunk_size (int | None): Number of output points per chunk.

Neighbor truncation

Neighbors beyond the max_neighbors nearest are dropped silently. With a weighting_fn a dropped neighbor can have a nonzero weight, which reintroduces a discontinuity in x. Choose max_neighbors large enough that truncation does not occur if you rely on derivatives of the output.

kernel_in_channels property

Number of input channels of the kernel network.

__call__(y: Float[Array, 'n d'], x: Float[Array, 'm d'], f_y: Float[Array, 'in_c n'] | None = None, *, quadrature_weights: Float[Array, n] | None = None, neighbors: neojax.nn.neighbor_search.Neighbors | None = None) -> Float[Array, 'out_c m']

Computes the neighbor search and kernel integral transform.

Parameters:

Name	Type	Description	Default
y	Float[Array, 'n d']	Coordinates of the n input points to integrate over.	required
x	Float[Array, 'm d']	Coordinates of the m output points. Must share the domain with y.	required
f_y	Float[Array, 'in_c n'] | None	Input function sampled at y. See IntegralTransform for when it is required.	None
quadrature_weights	Float[Array, n] | None	Optional quadrature weights of the n input points, see IntegralTransform. Default is None.	None
neighbors	Neighbors | None	Optional precomputed neighbors in y of each point in x, e.g. from radius_neighbor_search. Skips the neighbor search. Useful if the geometry is fixed. Default is None.	None

Returns:

Type	Description
Float[Array, 'out_c m']	The output function sampled at x.