Skip to content

Pointwise MLP

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 hidden layers with no final activation.

 gelu
use_bias bool

Whether to use bias. Default is True.

 True
dropout float

Dropout probability applied after each layer (except the last). If 0.0, no dropout is applied. Defaults to 0.0.

 0.0

Raises:

Type Description
ValueError

If dropout probability is non-float or not in the half-open interval [0,1).
ValueError

If too few or too many activation functions are passed.
Internal Attributes

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

  • weights (tuple[Float[Array, ...], ...]): Learnable weights.
  • biases (tuple[Float[Array, ...] | None, ...]): Learnable biases if use_bias is True.
  • activations (tuple[Callable, ...]): Activation functions between layers.
  • dropout (float): Dropout probability.

Examples:

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], dropout=0.1
)
x = jnp.ones((64, 32, 32))
out = mlp(x, key=key)

Tip

A single layer pointwise MLP is equivalent to a 1x1 convolution.

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

MLP forward pass.

Parameters:

Name Type Description Default
x Inexact[Array, 'in_c ...']

Input array.

 required
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 ...']

Output array.

Raises:

Type Description
ValueError

If the size of the leading channel axis does not match the MLP's input dimension.