Transolver Block

Transolver Block

neojax.nn.transolver_block.TransolverBlock

Transolver block.

A pre-norm transformer encoder block whose self-attention is replaced by PhysicsAttention, i.e. $$ f' = f + \mathrm{PhysicsAttention}(\mathrm{LN}(f)), \quad f'' = f' + \mathrm{MLP}(\mathrm{LN}(f')). $$

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key for parameter initialization.	required
hidden_channels	int	Number of input and output channels.	required
n_heads	int	Number of attention heads. Default is 8.	8
slice_num	int	Number of slices of the PhysicsAttention. Default is 32.	32
mlp_ratio	float	Ratio of the MLP hidden width to hidden_channels. Default is 4.	4.0
dropout	float	Dropout probability inside the PhysicsAttention. Default is 0.0.	0.0
activation	Callable	Activation function of the MLP. Default is jax.nn.gelu.	gelu
ndim	int | None	Number of spatial dimensions for structured grid inputs (convolutional attention projections) or None for pointwise projections (e.g. unstructured meshes). See PhysicsAttention. Default is None.	None
kernel_size	int	Kernel size of the convolutional attention projections. Default is 3.	3
temperature_range	tuple[float, float] | None	Clipping interval of the slice temperature. See PhysicsAttention. Default is (0.1, 5.0).	(0.1, 5.0)

Raises:

Type	Description
ValueError	If hidden_channels is not divisible by n_heads, or mlp_ratio is not positive.

Internal Attributes

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

• norm_1 (ChannelLayerNorm): Layer norm before the attention.
• attn (PhysicsAttention): The physics attention.
• norm_2 (ChannelLayerNorm): Layer norm before the MLP.
• mlp (PointwiseMLP): The pointwise MLP.

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

Forward pass of the Transolver block.

Parameters:

Name	Type	Description	Default
x	Inexact[Array, 'c ...']	Input array shaped (hidden_channels, N) or (hidden_channels, d1, ..., dN).	required
key	PRNGKeyArray | None	PRNG key used for dropout masks.	None
inference	bool	If True, dropout is disabled.	False

Returns:

Type	Description
Inexact[Array, 'c ...']	Output array with the same shape as x.