Components Reference¤

This page contains the API reference for the building blocks of Neural Operators in neojax.

While you can easily use pre-built models like FNO or TFNO, these underlying components are designed to be highly modular, allowing you to construct entirely custom neural operator architectures using standard equinox composition.

Spectral Convolution¤

These layers evaluate the continuous integral operator in Fourier space. They perform the core global operations that make Fourier Neural Operators discretization-invariant.

`neojax.nn.spectral_conv.SpectralConvNd` ¤

General n-dimensional spectral convolution layer.

This layer computes the real n-dimensional forward FFT, truncates the higher frequency modes according to the specified modes, multiplies the remaining modes with learnable complex weights, and transforms the result back to the spatial domain using the inverse FFT.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for weight initialization.	required
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int`	Number of output channels.	required
`modes`	`int \| Sequence[int]`	Number of Fourier modes to retain across each spatial dim.	required

Internal Attributes

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

weights (tuple[Complex[Array, ...], ...]): Learnable complex weights for the retained Fourier modes.
in_channels (int): Number of input channels.
out_channels (int): Number of output channels.
modes (tuple[int, ...]): Number of Fourier modes to retain across each spatial dim.

Example

import jax.numpy as jnp
import jax.random as jr
from neojax.nn import SpectralConvNd

key = jr.key(0)

# 1D Spectral Convolution
conv1d = SpectralConvNd(key, 3, 16, modes=16)
# or
conv1d = SpectralConvNd(key, 3, 16, modes=(16,))

# 2D Spectral Convolution
conv2d = SpectralConvNd(key, 3, 16, modes=(16, 16))

# 3D Spectral Convolution
conv3d = SpectralConvNd(key, 3, 16, modes=(16, 16, 16))

# Input shape: (channels, d1, ..., dN)
x = jnp.ones((3, 32, 32))
out = conv2d(x)

Upcoming Features

The current implementation doesn't support complex inputs (always truncates final dim currently). The addition is planned for an upcoming release.

`neojax.nn.tucker_spectral_conv.TuckerSpectralConvNd` ¤

General n-dimensional factorized spectral convolution layer.

This layer computes the real n-dimensional forward FFT, truncates the higher frequency modes according to the specified modes, multiplies the remaining modes with the Tucker factorization of the learnable complex weight tensors and transforms the result back to the spatial domain using the inverse FFT.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for weight initialization.	required
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int`	Number of output channels.	required
`modes`	`int \| Sequence[int]`	Number of Fourier modes to retain across each spatial dim.	required
`ranks`	`int \| Sequence[int]`	Number of ranks to contract the spectral tensors to. If `ranks` is an Integer, the same number is used for all ranks. If not, should be num_spatial_dims + num_channel_dims ranks, e.g. in 2D 4 ranks (out_channel, in_channel, 2 spatial ranks).	required
`share_factor_matrices`	`bool`	Whether to share the factor matrices across the weight tensors. This further decreases the number of weight tensors. Only the core tensor is not shared then. Default is True.	`True`

Example

import jax.numpy as jnp
import jax.random as jr
from neojax.nn import TuckerSpectralConvNd

key = jr.key(0)
x = jnp.ones((3, 32, 32))
tconv2d = TuckerSpectralConvNd(
    key, 3, 16, modes=(16, 16), ranks=4
)
out = tconv2d(x)

Upcoming Features

The current implementation doesn't support complex inputs (always truncates final dim currently). The addition is planned for an upcoming release.

Internal Attributes

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

core_tensors (tuple[Complex[Array, ...], ...]): Learnable complex core weight tensors.
factor_matrices (tuple[Complex[Array, ...], ...]): Learnable factor matrices for each dimension.
einsum_str (str): Static contraction string passed directly to jnp.einsum.
ndim (int): Total dimensions of the tensor structure (channels + spatial).

Cite

Useful overview on Tensor Operations. Tensor Decompositions and Applications

@article{kolda2009tensor,
    title={Tensor decompositions and applications},
    author={Kolda, Tamara G and Bader, Brett W},
    journal={SIAM review},
    volume={51},
    number={3},
    pages={455--500},
    year={2009},
    publisher={SIAM}
}

FNO Blocks¤

The standard layer of the Fourier Neural Operator. An FNO Block computes the sum of the global spectral convolution and a local skip connection, followed by optional normalization and non-linear activation.

`neojax.nn.fno_blocks.FNOBlocks` ¤

General FNO Blocks with variable layer number.

Implemented as in [1] and [2]. Each block is an instance of FNOBlock with an optional channelwise MLP.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for parameter initialization.	required
`n_layers`	`int`	Number of consecutive FNO blocks.	required
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int`	Number of output channels.	required
`modes`	`int \| Sequence[int]`	Number of Fourier modes to retain across each spatial dimension, e.g. (16, 16) for 2D inputs.	required
`activation`	`Callable \| Sequence[Callable]`	Activation function or sequence of activation functions used inside the FNO blocks. Default is `jax.nn.gelu`.	`gelu`
`use_channel_mlp`	`bool`	Whether to apply a pointwise channel MLP with 2 layers after each FNO block. Default is True.	`True`
`preactivation`	`bool`	Whether to apply the activation function before the spectral convolution and skip connections. Default is False. This is an exclusive flag. If activation is applied before, it isn't applied after.	`False`
`normalization`	`Literal['layer', 'instance', 'group'] \| None`	Type of normalization to use. Applied after the spectral_op + local_op sum, before activation. Can be `"layer"`, `"instance"`, `"group"` or None. Default is `"layer"`.	`'layer'`
`norm_groups`	`int`	Number of groups to use if `normalization="group"`. Defaults to 1.	`1`
`use_fno_residual`	`bool`	Whether to use a Resnet-style residual connection around each FNO block. Improves stability. Default True.	`True`
`local_operator`	`Literal['linear', 'soft-gating', 'identity'] \| None`	Type of local operator to use in the FNO block. 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`	`Literal['linear', 'soft-gating', 'identity'] \| None`	Type of residual connection used around the channel MLPs. Can be `"linear"`, `"soft-gating"`, `"identity"`, or None. Default is `"identity"`, which corresponds to Resnet-style residual connection.	`'identity'`
`channel_mlp_expansion`	`float \| None`	Expansion factor for computing the hidden channel dimension of the 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`

Internal Attributes

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

fno_layers (tuple[FNOBlock, ...]): The initialized FNOBlock layers.
channel_mlps (tuple[PointwiseMLP, ...] | None): The initialized PointwiseMLP layers, or None if use_channel_mlp is False.
channel_mlp_residuals (tuple | None): The residual connection instances or None.

Examples:

import jax.random as jr
import jax.numpy as jnp
from neojax.nn.fno_blocks import FNOBlocks

key = jr.key(0)

# Initialize a sequence of 4 FNO Blocks
fno_blocks = FNOBlocks(
    key=key,
    n_layers=4,
    in_channels=3,
    out_channels=16,
    modes=(8, 8),
    use_channel_mlp=True
)
# Input shape: (channels, height, width)
x = jnp.ones((3, 64, 64))
out = fno_blocks(x)

Cite

Fourier Neural Operator for Parametric Partial Differential Equations

@inproceedings{
    li2021fourier,
    title={Fourier Neural Operator for Parametric Partial Differential Equations},
    author={Zongyi Li and Nikola Kovachki and
    Kamyar Azizzadenesheli and Burigede liu and
    Kaushik Bhattacharya and Andrew Stuart and Anima Anandkumar},
    booktitle={International Conference on Learning Representations},
    year={2021},
    url={https://openreview.net/forum?id=c8P9NQVtmnO}
}

Neural Operator: Learning Maps Between Function Spaces With Applications to PDEs

@article{kovachki2023neural,
    title={Neural operator: Learning maps between function spaces with applications to pdes},
    author={Kovachki, Nikola and Li, Zongyi and
    Liu, Burigede and Azizzadenesheli,
    Kamyar and Bhattacharya, Kaushik and Stuart, Andrew and Anandkumar, Anima},
    journal={Journal of Machine Learning Research},
    volume={24},
    number={89},
    pages={1--97},
    year={2023}
}

Upcoming Features

The current implementation doesn't support dropout for the channel-wise MLP or normalization layers. Both will be added in future releases.

`neojax.nn.fno_blocks.FNOBlock` ¤

General FNO Block.

FNO block composed of the global non-linear integral kernel, i.e., \(\mathcal{K}(\mathbf{v}) = \mathcal{F}^{-1}(R \cdot \mathcal{F}(\mathbf{v}))\), and an optional local operator \(W \mathbf{v}\). A block can be further customized with Resnet-style residual connections around the FNO and normalization before the non-linear activation.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for parameter initialization.	required
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int`	Number of output channels.	required
`modes`	`int \| Sequence[int]`	Number of Fourier modes to retain across each spatial dimension. Must be a sequence of integers (e.g., `(16, 16)` for 2D) or single integer for 1D inputs.	required
`activation`	`Callable`	Activation function to use. Defaults to `jax.nn.gelu`.	`gelu`
`local_operator`	`Literal['linear', 'soft-gating', 'identity'] \| None`	Type of local operator to use. 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. Defaults to `False`.	`False`
`normalization`	`Literal['layer', 'instance', 'group'] \| None`	Type of normalization to use. Applied after the spectral_op + local_op sum, before activation. Can be `"layer"`, `"instance"`, `"group"` or None. Default is `"layer"`.	`'layer'`
`use_fno_residual`	`bool`	Whether to use a Resnet-style residual connection around each FNO block. Improves stability. Default True.	`True`
`preactivation`	`bool`	Whether to apply the activation before the spectral convolution and skip connection. Defaults to `False`. This is an exclusive flag. If activation is applied before, it isn't applied after.	`False`

Internal Attributes

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

spectral_conv (SpectralConvNd): The SpectralConvNd layer performing the operator integral.
local_operator (Flattened1dConv | SoftGating | eqx.nn.Identity | None): The initialized local operator layer or None.
normalization (InstanceNorm | eqx.nn.GroupNorm | None): Type of normalization to use. Applied after the spectral_op + local_op sum, before activation.
activation (Callable): The activation function.
preactivation (bool): Boolean flag indicating if preactivation is used.
use_fno_residual (bool): Whether to use a Resnet-style residual connection around each FNO block. Improves stability.

Examples:

import jax.random as jr
from neojax.nn.fno_blocks import FNOBlock
import jax.numpy as jnp

key = jr.key(0)
# Initialize a 2D FNO Block
fno_block = FNOBlock(
    key=key,
    in_channels=3,
    out_channels=8,
    modes=(4, 4),
    local_operator="linear"
)
# Input shape: (channels, height, width)
x = jnp.ones((3, 32, 32))
out = fno_block(x)

TFNO Blocks¤

Tucker-factorized FNO blocks. By factorizing the spectral weights into a core tensor and factor matrices, TFNO blocks significantly reduce the parameter count and memory footprint, especially for 3D or 4D problems.

`neojax.nn.tfno_blocks.TFNOBlocks` ¤

General TFNO Blocks with variable layer number.

Each block is an instance of TFNOBlock with an optional channelwise MLP.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for parameter initialization.	required
`n_layers`	`int`	Number of consecutive TFNO blocks.	required
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int`	Number of output channels.	required
`modes`	`int \| Sequence[int]`	Number of Fourier modes to retain across each spatial dimension, e.g. (16, 16) for 2D inputs.	required
`ranks`	`int \| Sequence[int]`	Number of ranks to contract the spectral tensors to. If `ranks` is an Integer, the same number is used for all ranks. If not, should be num_spatial_dims + num_channel_dims ranks, e.g. in 2D 4 ranks (out_channel, in_channel, 2 spatial ranks).	required
`share_factor_matrices`	`bool`	Whether to share the factor matrices across the weight tensors. This further decreases the number of weight tensors. Only the core tensor is not shared then. Default is True.	`True`
`activation`	`Callable \| Sequence[Callable]`	Activation function or sequence of activation functions used inside the TFNO blocks. Default is `jax.nn.gelu`.	`gelu`
`use_channel_mlp`	`bool`	Whether to apply a pointwise channel MLP with 2 layers after each TFNO block. Default is True.	`True`
`preactivation`	`bool`	Whether to apply the activation function before the spectral convolution and skip connections. Default is False. This is an exclusive flag. If activation is applied before, it isn't applied after.	`False`
`normalization`	`Literal['layer', 'instance', 'group'] \| None`	Type of normalization to use. Applied after the spectral_op + local_op sum, before activation. Can be `"layer"`, `"instance"`, `"group"` or None. Default is `"layer"`.	`'layer'`
`norm_groups`	`int`	Number of groups to use if `normalization="group"`. Defaults to 1.	`1`
`use_fno_residual`	`bool`	Whether to use a Resnet-style residual connection around each TFNO block. Improves stability. Default True.	`True`
`local_operator`	`Literal['linear', 'soft-gating', 'identity'] \| None`	Type of local operator to use in the TFNO block. 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 TFNO blocks local operator. Default is `False`.	`False`
`channel_mlp_residual`	`Literal['linear', 'soft-gating', 'identity'] \| None`	Type of residual connection used around the channel MLPs. Can be `"linear"`, `"soft-gating"`, `"identity"`, or None. Default is `"identity"`, which corresponds to Resnet-style residual connection.	`'identity'`
`channel_mlp_expansion`	`float \| None`	Expansion factor for computing the hidden channel dimension of the 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`

Internal Attributes

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

tfno_layers (tuple[TFNOBlock, ...]): The initialized TFNOBlock layers.
channel_mlps (tuple[PointwiseMLP, ...] | None): The initialized PointwiseMLP layers, or None if use_channel_mlp is False.
channel_mlp_residuals (tuple | None): The residual connection instances or None.

Examples:

import jax.random as jr
import jax.numpy as jnp
from neojax.nn.tfno_blocks import TFNOBlocks

key = jr.key(0)

# Initialize a sequence of 4 TFNO Blocks
tfno_blocks = TFNOBlocks(
    key=key,
    n_layers=4,
    in_channels=3,
    out_channels=16,
    modes=(8, 8),
    use_channel_mlp=True
)
# Input shape: (channels, height, width)
x = jnp.ones((3, 64, 64))
out = tfno_blocks(x)

Upcoming Features

The current implementation doesn't support dropout for the channel-wise MLP or normalization layers. Both will be added in future releases.

`neojax.nn.tfno_blocks.TFNOBlock` ¤

General Tucker factorized FNO Block.

TFNO block composed of the global non-linear integral kernel, i.e., \(\mathcal{K}(\mathbf{v}) = \mathcal{F}^{-1}(R \cdot \mathcal{F}(\mathbf{v}))\), and an optional local operator \(W \mathbf{v}\). Here \(R\) is decomposed into core tensor and factor matrices. A block can be further customized with Resnet-style residual connections around the TFNO and normalization before the non-linear activation.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for parameter initialization.	required
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int`	Number of output channels.	required
`modes`	`int \| Sequence[int]`	Number of Fourier modes to retain across each spatial dimension. Must be a sequence of integers (e.g., `(16, 16)` for 2D) or single integer for 1D inputs.	required
`ranks`	`int \| Sequence[int]`	Number of ranks to contract the spectral tensors to. If `ranks` is an Integer, the same number is used for all ranks. If not, should be num_spatial_dims + num_channel_dims ranks, e.g. in 2D 4 ranks (out_channel, in_channel, 2 spatial ranks).	required
`share_factor_matrices`	`bool`	Whether to share the factor matrices across the weight tensors. This further decreases the number of weight tensors. Only the core tensor is not shared then. Default is True.	`True`
`activation`	`Callable`	Activation function to use. Defaults to `jax.nn.gelu`.	`gelu`
`local_operator`	`Literal['linear', 'soft-gating', 'identity'] \| None`	Type of local operator to use. 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 TFNO blocks local operator. Defaults to `False`.	`False`
`normalization`	`Literal['layer', 'instance', 'group'] \| None`	Type of normalization to use. Applied after the spectral_op + local_op sum, before activation. Can be `"layer"`, `"instance"`, `"group"` or None. Default is `"layer"`.	`'layer'`
`use_fno_residual`	`bool`	Whether to use a Resnet-style residual connection around each TFNO block. Improves stability. Default True.	`True`
`preactivation`	`bool`	Whether to apply the activation before the spectral convolution and skip connection. Defaults to `False`. This is an exclusive flag. If activation is applied before, it isn't applied after.	`False`

Internal Attributes

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

tucker_spectral_conv (TuckerSpectralConvNd): The TuckerSpectralConvNd layer performing the operator integral.
local_operator (Flattened1dConv | SoftGating | eqx.nn.Identity | None): The initialized local operator layer or None.
normalization (InstanceNorm | eqx.nn.GroupNorm | None): Type of normalization to use. Applied after the spectral_op + local_op sum, before activation.
activation (Callable): The activation function.
preactivation (bool): Boolean flag indicating if preactivation is used.
use_fno_residual (bool): Whether to use a Resnet-style residual connection around each TFNO block. Improves stability.

Examples:

import jax.random as jr
from neojax.nn.tfno_blocks import TFNOBlock
import jax.numpy as jnp

key = jr.key(0)
# Initialize a 2D TFNO Block
tfno_block = TFNOBlock(
    key=key,
    in_channels=3,
    out_channels=8,
    modes=(4, 4),
    ranks=2,
    local_operator="linear"
)
# Input shape: (channels, height, width)
x = jnp.ones((3, 32, 32))
out = tfno_block(x)

Pointwise MLP¤

Applies a Multi-Layer Perceptron independently across the spatial grid points, operating solely on the channel dimension. These are used for lifting inputs to higher-dimensional latent spaces, projecting outputs, and channel-mixing within operator blocks.

`neojax.nn.pointwise_mlp.PointwiseMLP` ¤

General pointwise MLP.

Parameters:

Name	Type	Description	Default
`key`	`PRNGKeyArray`	PRNG key for weight initialization.	required
`layers`	`Sequence[int]`	Sequence of dimensions for each layer.	required
`activations`	`Callable \| Sequence[Callable]`	Sequence of activations between layers or single activation for all layers. If single activation, MLP is initialized with given l layers and l - 1 instances of type `activation`. Otherwise, a final activation after the last layer can be passed. Default is GeLu activation for all layers.	`gelu`

Internal Attributes

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

weights (tuple[Float[Array, ...], ...]): Learnable weights.
biases (tuple[Float[Array, ...], ...]): Learnable biases.
activations (tuple[Callable, ...]): Activation functions between layers.

Example

import jax.numpy as jnp
import jax.random as jr
from neojax.nn import PointwiseMLP

key = jr.key(0)

mlp = PointwiseMLP(
    key, [64, 128, 128], [jax.nn.gelu, jax.nn.gelu]
)
x = jnp.ones((64, 32, 32))
out = mlp(x)

Domain Padding¤

Fast Fourier Transforms (FFT) assume periodic boundary conditions. When learning on non-periodic domains, DomainPadding pads the domain before the spectral convolutions and unpads it afterward, severely mitigating boundary artifacts.

`neojax.nn.domain_padding.DomainPadding` ¤

Applies domain padding to a signal.

Domain Padding helps to mitigate boundary artifacts. In the context of FNO, it is often used to handle non-periodic boundary conditions by padding the spatial dimensions before the spectral convolution and cropping the result back to the original resolution. Padding is applied symmetrically to each dimension, i.e., a 20% (0.2) padding pads an input by 10% on each side along an axis.

Parameters:

Name	Type	Description	Default
`padding`	`float \| Sequence[float]`	Percentage of padding to apply to each spatial dimension. Can be a single float in [0, 1] (applied to all dims) or a sequence of floats.	required
`mode`	`str`	The type of padding to apply (e.g. "constant", "edge"). Defaults to "constant" which pads with zeros. See https://docs.jax.dev/en/latest/_autosummary/jax.numpy.pad.html#jax.numpy.pad for all padding modes.	`'constant'`

Internal Attributes

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

padding (float | Sequence[float]): The stored padding ratios for each dimension.
mode (str): The padding mode.

Info

Currently doesn't support resolution scaling.

`mode` `class-attribute` ¤

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

`pad(x: Float[Array, 'c ...']) -> Float[Array, 'c ...']` ¤

Pads the input array based on the configured ratios.

Parameters:

Name	Type	Description	Default
`x`	`Float[Array, 'c ...']`	The input array of shape `(channels, d1, ..., dn)`.	required

Returns:

Type	Description
`Float[Array, 'c ...']`	The padded array.

`unpad(x: Float[Array, 'c ...'], original_shape: tuple[int, ...]) -> Float[Array, 'c ...']` ¤

Crops the padded array back to its original resolution.

Parameters:

Name	Type	Description	Default
`x`	`Float[Array, 'c ...']`	The padded array.	required
`original_shape`	`tuple[int, ...]`	The original array shape.	required

Returns:

Type	Description
`Float[Array, 'c ...']`	The cropped array of original resolution.

Positional Embedding¤

Appends grid coordinate features (e.g., \((x, y)\) positions) to the input tensors channel dimension.

`neojax.nn.positional_embedding.GridEmbeddingNd` ¤

Regular grid embedding for n-dim signals.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Number of input channels.	required
`ndim`	`int`	Number of spatial dimensions.	required
`grid_boundaries`	`tuple[tuple[int, int], ...] \| None`	Boundaries of regular grid per dimension ((low, high), ...) or None. Default is None which is a ((0, 1), ...) bounded grid.	`None`

Internal Attributes

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

in_channels (int): Number of input channels.
grid_boundaries (tuple[tuple[int, int], ...]): Boundaries of regular grid.

Skip Connections¤

Local operators used alongside the global spectral convolutions. They process high-frequency, localized information and act as residual connections to stabilize training.

Note: For the standard identity skip connection, neojax directly uses equinox.nn.Identity for simplicity and seamless integration with the JAX/Equinox ecosystem.

`neojax.nn.skip_connections.SoftGating` ¤

Applies soft-gating by weighting the channels of the given input.

Given an input x of size (batch-size, channels, height, width), this returns x * w where w is of shape (1, channels, 1, 1)

Parameters:

Name	Type	Description	Default
`ndim`	`int`	Number of spatial dimensions.	required
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int \| None`	Number of output channels. If provided, must match `in_channels`.	`None`
`use_bias`	`bool`	Whether to include a learnable bias.	`False`

Internal Attributes

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

weight (Float[Array, ...]): Learnable channel-wise weights.
bias (Float[Array, ...] | None): Optional learnable channel-wise bias.

`neojax.nn.skip_connections.Flattened1dConv` ¤

Applies 1d convolution to flattened dimensions.

Skip layer that flattens all dimensions except for the leading channel dimension and applies 1d convolution over them, then un-flattens result.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Number of input channels.	required
`out_channels`	`int`	Number of output channels.	required
`kernel_size`	`int`	Size of the convolving kernel.	required
`use_bias`	`bool`	Whether to add a learnable bias to the output.	`False`

Internal Attributes

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

conv (eqx.nn.Conv1d): The underlying 1D convolution layer.
out_channels (int): Number of output channels.

Components Reference¤

Spectral Convolution¤

neojax.nn.spectral_conv.SpectralConvNd ¤

neojax.nn.tucker_spectral_conv.TuckerSpectralConvNd ¤

FNO Blocks¤

neojax.nn.fno_blocks.FNOBlocks ¤

neojax.nn.fno_blocks.FNOBlock ¤

TFNO Blocks¤

neojax.nn.tfno_blocks.TFNOBlocks ¤

neojax.nn.tfno_blocks.TFNOBlock ¤

Pointwise MLP¤

neojax.nn.pointwise_mlp.PointwiseMLP ¤

Domain Padding¤

neojax.nn.domain_padding.DomainPadding ¤

mode class-attribute ¤

pad(x: Float[Array, 'c ...']) -> Float[Array, 'c ...'] ¤

unpad(x: Float[Array, 'c ...'], original_shape: tuple[int, ...]) -> Float[Array, 'c ...'] ¤

Positional Embedding¤

neojax.nn.positional_embedding.GridEmbeddingNd ¤

Skip Connections¤

neojax.nn.skip_connections.SoftGating ¤

neojax.nn.skip_connections.Flattened1dConv ¤

`neojax.nn.spectral_conv.SpectralConvNd` ¤

`neojax.nn.tucker_spectral_conv.TuckerSpectralConvNd` ¤

`neojax.nn.fno_blocks.FNOBlocks` ¤

`neojax.nn.fno_blocks.FNOBlock` ¤

`neojax.nn.tfno_blocks.TFNOBlocks` ¤

`neojax.nn.tfno_blocks.TFNOBlock` ¤

`neojax.nn.pointwise_mlp.PointwiseMLP` ¤

`neojax.nn.domain_padding.DomainPadding` ¤

`mode` `class-attribute` ¤

`pad(x: Float[Array, 'c ...']) -> Float[Array, 'c ...']` ¤

`unpad(x: Float[Array, 'c ...'], original_shape: tuple[int, ...]) -> Float[Array, 'c ...']` ¤

`neojax.nn.positional_embedding.GridEmbeddingNd` ¤

`neojax.nn.skip_connections.SoftGating` ¤

`neojax.nn.skip_connections.Flattened1dConv` ¤