Normalization

Instance Normalization

neojax.nn.normalization.InstanceNorm

Computes per instance normalization.

Computes $$ \hat{x_{b,c,\mathbf{d}}} = \frac{x_{b,c,\mathbf{d}} - \mu_{b,c}}{\sqrt{\sigma_{b,c}^2 + \epsilon}} $$ where the mean and variance are over the spatial dimensions only. Optionally applies the transform \(\gamma_c \hat{x}_{b,c,\mathbf{d}} + \beta_c\) with learnable parameters \(\gamma_c\) and \(\beta_c\).

Parameters:

Name	Type	Description	Default
shape	Sequence[int]	Overall shape of the input signal, i.e., (channels, d1, ..., dN).	required
eps	float	A small value added to the variance for numerical stability. Defaults to 1e-7.	1e-07
use_weight	bool	Whether to include a learnable affine weight. Defaults to True.	True
use_bias	bool	Whether to include a learnable affine bias. Defaults to True.	True

Internal Attributes

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

• shape (tuple[int, ...]): The input shape.
• eps (float): The numerical stability value.
• weight (Float[Array, ...] | None): The learnable weights or None.
• bias (Float[Array, ...] | None): The learnable bias or None.
• use_weight (bool): Whether to include a learnable affine weight.
• use_bias (bool): Whether to include a learnable affine bias.

__call__(x: Inexact[Array, 'c ...']) -> Inexact[Array, 'c ...']

Compute Instance Norm.

Parameters:

Name	Type	Description	Default
x	Inexact[Array, 'c ...']	Input array with the same shape as shape in __init__ (channel, d1, ..., dN).	required

Returns:

Type	Description
Inexact[Array, 'c ...']	Normalized array.

---

neojax.nn.normalization.ChannelLayerNorm

Layer normalization over the channel axis of a channel-first signal.

Wraps eqx.nn.LayerNorm, which normalizes the trailing axis of an unbatched input, and applies it independently at every spatial location of a channel-first array, i.e.

\[ \hat{x}_{c,\mathbf{d}} = \gamma_c \frac{x_{c,\mathbf{d}} - \mu_{\mathbf{d}}}{\sqrt{\sigma_{\mathbf{d}}^2 + \epsilon}} + \beta_c, \]

where the mean and variance are over the channel axis only.

Parameters:

Name	Type	Description	Default
channels	int	Number of channels of the input signal.	required
eps	float	A small value added to the variance for numerical stability. Defaults to 1e-5.	1e-05
use_weight	bool	Whether to include a learnable affine weight. Defaults to True.	True
use_bias	bool	Whether to include a learnable affine bias. Defaults to True.	True

Internal Attributes

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

• norm (eqx.nn.LayerNorm): The wrapped layer normalization.

__call__(x: Inexact[Array, 'c ...']) -> Inexact[Array, 'c ...']

Compute layer norm over the channel axis.

Parameters:

Name	Type	Description	Default
x	Inexact[Array, 'c ...']	Input array shaped (channels, d1, ..., dN).	required

Returns:

Type	Description
Inexact[Array, 'c ...']	Normalized array of the same shape.