Transolver

Transolver

neojax.models.transolver.Transolver

Transolver Neural Operator.

The model consists of a lifting MLP that maps the input features (and positional information) to a latent space, a learnable placeholder that is added to every point, a stack of TransolverBlocks, and a final layer norm and linear projection. The implementation follows the reference.

The model supports unstructured meshes (use_mesh=True), where an input signal u of shape (in_channels, N) is accompanied by the coordinates x_in of its N nodes, as well as structured grids (use_mesh=False) of shape (in_channels, d1, ..., d_ndim), where the attention projections are convolutions.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key for parameter initialization.	required
in_channels	int	Number of input channels of the signal u. On structured grids, this includes the coordinates if they are part of the signal (see positional_embedding).	required
out_channels	int	Number of output channels.	required
ndim	int	Number of spatial dimensions (dimension of the coordinates).	required
hidden_channels	int	Hidden channel dimension (latent width). Default is 256.	256
n_layers	int	Number of Transolver blocks. Default is 5.	5
n_heads	int	Number of attention heads. hidden_channels must be divisible by n_heads. Default is 8.	8
slice_num	int	Number of slices (slice tokens) of the physics attention. Default is 32.	32
mlp_ratio	float	Ratio of the hidden width of the MLP in each block to hidden_channels. Default is 1.0.	1.0
dropout	float	Dropout probability in the physics attention. Default is 0.0.	0.0
activation	Callable	Activation function used within the layers. Defaults to jax.nn.gelu.	gelu
use_mesh	bool	Whether the model operates on unstructured meshes (True) or structured grids (False). On meshes, the node coordinates x_in are required in the forward pass and are appended to the input signal and the attention uses pointwise projections. On grids, the attention uses convolutional projections. Default is False.	False
kernel_size	int	Kernel size of the convolutional attention projections on structured grids. Ignored if use_mesh is True. Must be odd. Default is 3.	3
lift_channel_ratio	float	Ratio of the hidden width of the lifting MLP to hidden_channels. Default is 2.0.	2.0
positional_embedding	Union[GridEmbeddingNd, Literal['grid'], NoneType]	Positional embedding to apply to the last channels of the raw input on structured grids. Pass "grid" for a regular GridEmbeddingNd on a ((0, 1), ...) grid, a GridEmbeddingNd instance, or None if the coordinates are part of the input signal already (e.g., by using ConcatenateCoordsSchema) or not used. Not supported if use_mesh is True (node coordinates are always used) or unified_pos is True. Default is None.	None
unified_pos	bool	Whether to use the "unified position" embedding of the reference: instead of the raw coordinates, the distances of every point to the nodes of a regular ref ** ndim reference grid (see ReferenceGridEmbedding) are appended to the input signal. Default is False.	False
ref	int	Number of reference grid nodes per dimension of the unified position embedding. Default is 8.	8
temperature_range	tuple[float, float] | None	Interval the learnable slice temperatures are clipped to in the forward pass for stability. None disables clipping. Default is (0.1, 5.0).	(0.1, 5.0)
init_std	float | None	Standard deviation of the truncated normal initialization of all linear weights, as in the reference. Note that this also overrides the orthogonal initialization of the slice projections (as the reference does). None keeps the default initialization of each component. Default is 0.02.	0.02

Raises:

Type	Description
ValueError	If any of in_channels, out_channels, ndim, hidden_channels, n_layers, ref is not a positive integer.
ValueError	If positional_embedding is invalid, is passed with use_mesh=True or together with unified_pos=True, or does not match ndim.
ValueError	If init_std is not positive.

Internal Attributes

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

• positional_embedding (GridEmbeddingNd | None): Grid positional embedding appended to the raw input signal.
• reference_embedding (ReferenceGridEmbeddingNd | None): Unified position embedding appended to the raw input signal.
• lifting (PointwiseMLP): Lifts the input (and positional features) to hidden_channels.
• placeholder (Float[Array, "c"]): Learnable vector added to every lifted point.
• blocks (tuple[TransolverBlock, ...]): The Transolver blocks.
• final_norm (ChannelLayerNorm): Layer norm before the projection.
• projection (PointwiseMLP): Linear map from hidden_channels to out_channels.
• ndim (int): Number of spatial dimensions.
• use_mesh (bool): Whether the model operates on unstructured meshes.

Cite

Transolver: A Fast Transformer Solver for PDEs on General Geometries

@inproceedings{wu2024transolver,
    title={Transolver: A Fast Transformer Solver for PDEs on General Geometries},
    author={Haixu Wu and Huakun Luo and Haowen Wang and
    Jianmin Wang and Mingsheng Long},
    booktitle={International Conference on Machine Learning},
    year={2024}
}

Examples:

import jax.numpy as jnp
import jax.random as jr
from neojax.models import Transolver

key = jr.key(0)

# unstructured mesh with 500 nodes in 2D
model = Transolver(
    key, in_channels=1, out_channels=1, ndim=2,
    hidden_channels=64, n_layers=4, use_mesh=True,
)
u, x_in = jnp.ones((1, 500)), jnp.ones((500, 2))
out = model(u, x_in=x_in)  # (1, 500)

# structured 2D grid with appended grid coordinates
model = Transolver(
    key, in_channels=1, out_channels=1, ndim=2,
    hidden_channels=64, n_layers=4, positional_embedding="grid",
)
out = model(jnp.ones((1, 32, 32)))  # (1, 32, 32)

No time embedding supported yet

neojax doesn't support the time embedding of Transolver yet.

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

Forward pass of the Transolver model.

Parameters:

Name	Type	Description	Default
u	Inexact[Array, 'in_c ...']	Input signal. Shape is (in_channels, N) if use_mesh is True or (in_channels, d1, ..., d_ndim) otherwise.	required
x_in	Float[Array, 'N d'] | None	Coordinates of the N nodes of u with shape (N, ndim). Required if use_mesh is True, must be None otherwise.	None
key	PRNGKeyArray | None	PRNG key used for dropout masks.	None
inference	bool	If True, dropout is disabled.	False

Returns:

Type	Description
Inexact[Array, 'out_c ...']	The predicted field of shape (out_channels, N) if use_mesh is True or (out_channels, d1, ..., d_ndim) otherwise.

Raises:

Type	Description
ValueError	If use_mesh is True and x_in is None or if u and x_in have mismatching shapes.
ValueError	If use_mesh is False and x_in is passed or u doesn't have ndim spatial axes.