Normalization

Normalization layers for deep neural networks.

class taktiny.nn.BatchNorm(num_features, epsilon=1e-06, *, bias=False, dtype=None, rngs=None, elementwise_affine=True, momentum=0.1, track_running_stats=True, axes=None, scale_initializer=<function ones>, bias_initializer=<function zeros>, quant=None, axis_names=None, partition_spec=None, scale_metadata=None, bias_metadata=None)

Bases: Module

Normalize feature dimensions using mini-batch statistics.

BatchNorm retains the dimensions selected by axes and computes a mean and variance by reducing every other input dimension. Its transformation is

output = (x - mean) / sqrt(variance + epsilon) * scale + bias.

num_features may contain one or more dimensions. Their sizes and order must match axes and determine the shapes of the affine parameters and running statistics. If axes is omitted, the final len(num_features) input dimensions are treated as features.

Mode is controlled by Module.is_training, which is changed through Module.train() and Module.eval(). During training, current batch statistics normalize the input. If track_running_stats=True, the same call also updates running_mean, running_var, and num_batches_tracked. During evaluation, the stored statistics are used. When tracking is disabled, current input statistics are used in both modes.

Running-state mutation is an eager side effect and is therefore rejected while tracing with jax.jit. Use evaluation mode or set track_running_stats=False for a directly jitted call. Float16 and bfloat16 statistics are accumulated in float32 before the output is cast back to its input dtype.

Parameters:

• num_features (int | Sequence[int]) – Size of every retained feature dimension. An integer is treated as a one-dimensional feature shape.

• epsilon (float) – Positive value added to the variance before the reciprocal square root. Defaults to 1e-6.

• bias (bool) – Whether to create an additive bias when affine parameters are enabled. Defaults to False.

• dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Dtype used for affine parameters and running statistics. Defaults to float32 through the default initializers.

• rngs (Rngs | None) – Random stream used by parameter initializers. A deterministic stream seeded with zero is used when omitted.

• elementwise_affine (bool) – Whether to create a learnable scale and optional bias with shape num_features. Defaults to True.

• momentum (float | None) – Weight assigned to the newest batch statistics when updating running state. None selects a cumulative moving average. Defaults to 0.1.

• track_running_stats (bool) – Whether to maintain statistics for evaluation. Defaults to True.

• axes (int | Sequence[int] | None) – Input feature dimensions retained by the statistics, in the same order as num_features. Every remaining dimension is reduced. Defaults to the trailing feature dimensions.

• scale_initializer (Initializer) – Initializer for the multiplicative scale. Defaults to ones.

• bias_initializer (Initializer) – Initializer for the additive bias. Defaults to zeros.

• quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix configuration. Rules for unrelated operation types are ignored; quantized normalization parameters are not yet supported.

• axis_names (tuple[str | None, ...] | None) – Logical names for all feature and running-state dimensions.

• partition_spec (P | None) – Explicit partition specification for affine parameters and running statistics. A mapping obtained from axis_names overrides this specification.

• scale_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Metadata attached to the scale parameter.

• bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Metadata attached to the bias parameter.

Variables:

• num_features – Retained feature shape as a tuple.

• axes – Input dimensions treated as features.

• scale – Learnable multiplicative parameter, or None when affine parameters are disabled.

• bias – Learnable additive parameter, or None when disabled.

• running_mean – Non-trainable running mean, or None when tracking is disabled.

• running_var – Non-trainable running variance, or None when tracking is disabled.

• num_batches_tracked – Number of eager training batches incorporated into the running statistics, or None when tracking is disabled.

Example

Train once to record statistics, then switch to evaluation mode:

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> layer = nn.BatchNorm(4, momentum=1.0)
>>> x = jnp.arange(24, dtype=jnp.float32).reshape(2, 3, 4)
>>> training_output = layer(x)
>>> _ = layer.eval()
>>> evaluation_output = layer(x)
>>> training_output.shape == evaluation_output.shape
True

References

Sergey Ioffe and Christian Szegedy, “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift” (2015). https://arxiv.org/abs/1502.03167

class taktiny.nn.GroupNorm(num_groups, num_channels, epsilon=1e-06, *, bias=False, dtype=None, rngs=None, elementwise_affine=True, channel_axes=-1, batch_axes=0, scale_initializer=<function ones>, bias_initializer=<function zeros>, quant=None, axis_names=None, partition_spec=None, scale_metadata=None, bias_metadata=None)

Bases: Module

Normalize grouped channels independently of the batch size.

GroupNorm divides each channel dimension into a corresponding number of groups. For every sample and Cartesian combination of channel groups, it computes mean and variance across the within-group channel dimensions and every spatial dimension:

output = (x - mean) / sqrt(variance + epsilon) * scale + bias.

num_channels and num_groups may both be N-D shapes. They must have equal rank, and every num_channels[i] must be divisible by num_groups[i]. Their dimensions correspond in order to channel_axes. With N-D channels and the default channel_axes=-1, the final len(num_channels) dimensions are selected automatically.

Dimensions in batch_axes are kept independent and never contribute to the statistics. The default is the leading dimension 0. Pass an empty tuple for unbatched input. Because GroupNorm never aggregates across batch dimensions and stores no running statistics, training and evaluation modes behave identically. Float16 and bfloat16 statistics are computed in float32, and the result is cast back to the input dtype.

Parameters:

• num_groups (int | Sequence[int]) – Number of groups used to split each channel dimension. An integer is treated as a one-dimensional group shape.

• num_channels (int | Sequence[int]) – Size of every channel dimension. It must have the same rank as num_groups.

• epsilon (float) – Positive value added to the variance before the reciprocal square root. Defaults to 1e-6.

• bias (bool) – Whether to create an additive bias when affine parameters are enabled. Defaults to False.

• dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Dtype passed to the parameter initializers. Defaults to each initializer’s default dtype.

• rngs (Rngs | None) – Random stream used by parameter initializers. A deterministic stream seeded with zero is used when omitted.

• elementwise_affine (bool) – Whether to create a learnable scale and optional bias with shape num_channels. Defaults to True.

• channel_axes (int | Sequence[int]) – Input dimensions corresponding to num_channels, in the same order. Defaults to the trailing channel dimensions.

• batch_axes (int | Sequence[int]) – Input dimensions kept independent during normalization. Defaults to the leading dimension 0. Use () for unbatched input.

• scale_initializer (Initializer) – Initializer for the multiplicative scale. Defaults to ones.

• bias_initializer (Initializer) – Initializer for the additive bias. Defaults to zeros.

• quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix configuration. Rules for unrelated operation types are ignored; quantized normalization parameters are not yet supported.

• axis_names (tuple[str | None, ...] | None) – Logical names for all num_channels dimensions.

• partition_spec (P | None) – Explicit partition specification for affine parameters. A mapping obtained from axis_names overrides this specification.

• scale_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Metadata attached to the scale parameter.

• bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Metadata attached to the bias parameter.

Variables:

• num_groups – Group shape as a tuple.

• num_channels – Channel shape as a tuple.

• channel_axes – Input dimensions containing channels.

• batch_axes – Input dimensions kept independent.

• scale – Learnable multiplicative parameter, or None when affine parameters are disabled.

• bias – Learnable additive parameter, or None when disabled.

Example

Normalize a tensor with two structured channel dimensions. The final dimensions (4, 6) are divided into (2, 2) groups:

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> layer = nn.GroupNorm((2, 2), (4, 6))
>>> x = jnp.ones((2, 5, 4, 6), dtype=jnp.float32)
>>> layer(x).shape
(2, 5, 4, 6)

References

Yuxin Wu and Kaiming He, “Group Normalization” (2018). https://arxiv.org/abs/1803.08494

class taktiny.nn.LayerNorm(num_features, epsilon=1e-06, *, bias=True, dtype=None, rngs=None, elementwise_affine=True, axes=None, scale_initializer=<function ones>, bias_initializer=<function zeros>, quant=None, axis_names=None, partition_spec=None, scale_metadata=None, bias_metadata=None)

Bases: Module

Apply layer normalization over selected feature dimensions.

For every independent slice of the input, LayerNorm computes the mean and variance over axes and applies

output = (x - mean) / sqrt(variance + epsilon) * scale + bias.

Unlike BatchNorm, its statistics do not depend on other samples and it performs the same calculation in training and evaluation modes. num_features declares the sizes of the normalized dimensions in the same order as axes. If axes is omitted, the final len(num_features) dimensions are normalized. For example, num_features=(3, 4) expects trailing dimensions of shape (3, 4).

Statistics for float16 and bfloat16 inputs are computed in float32 for numerical stability. The result is converted back to the input dtype.

Parameters:

• num_features (int | Sequence[int]) – Size of every normalized input dimension. An integer is treated as a one-dimensional feature shape.

• epsilon (float) – Positive value added to the variance before the reciprocal square root. Defaults to 1e-6.

• bias (bool) – Whether to create an additive bias when affine parameters are enabled. Defaults to True.

• dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Dtype passed to the parameter initializers. Defaults to each initializer’s default dtype.

• rngs (Rngs | None) – Random stream used by parameter initializers. A deterministic stream seeded with zero is used when omitted.

• elementwise_affine (bool) – Whether to create a learnable scale and optional bias with shape num_features. Defaults to True.

• axes (int | Sequence[int] | None) – Input dimensions corresponding to num_features, in the same order. Negative dimensions are supported. Defaults to the trailing feature dimensions.

• scale_initializer (Initializer) – Initializer for the multiplicative scale. Defaults to ones.

• bias_initializer (Initializer) – Initializer for the additive bias. Defaults to zeros.

• quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix configuration. Rules for unrelated operation types are ignored; quantized normalization parameters are not yet supported.

• axis_names (tuple[str | None, ...] | None) – Logical names for all num_features dimensions.

• partition_spec (P | None) – Explicit partition specification for affine parameters. A mapping obtained from axis_names overrides this specification.

• scale_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Metadata attached to the scale parameter.

• bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Metadata attached to the bias parameter.

Variables:

• num_features – Normalized feature shape as a tuple.

• axes – Input dimensions normalized by this module.

• scale – Learnable multiplicative parameter, or None when elementwise_affine=False.

• bias – Learnable additive parameter, or None when disabled.

Example

Normalize the final two dimensions independently for each item in the leading batch dimension:

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> layer = nn.LayerNorm((3, 4))
>>> x = jnp.arange(24, dtype=jnp.float32).reshape(2, 3, 4)
>>> layer(x).shape
(2, 3, 4)

References

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton, “Layer Normalization” (2016). https://arxiv.org/abs/1607.06450

class taktiny.nn.RMSNorm(num_features, epsilon=1e-06, *, bias=False, dtype=None, rngs=None, elementwise_affine=True, axes=None, scale_initializer=<function ones>, bias_initializer=<function zeros>, quant=None, axis_names=None, partition_spec=None, scale_metadata=None, bias_metadata=None)

Bases: Module

Apply root mean square normalization over selected dimensions.

RMSNorm scales each independent input slice using its root mean square:

output = x / sqrt(mean(x ** 2) + epsilon) * scale + bias.

It does not subtract the mean. Consequently, RMSNorm provides re-scaling invariance without the re-centering step performed by LayerNorm. The statistics are local to each input slice, so training and evaluation modes behave identically.

num_features gives the sizes of the normalized dimensions in the same order as axes. When axes is omitted, the final len(num_features) dimensions are used. Statistics for float16 and bfloat16 inputs are computed in float32, after which the result is cast back to the input dtype.

Parameters:

• num_features (int | Sequence[int]) – Size of every normalized input dimension. An integer is treated as a one-dimensional feature shape.

• epsilon (float) – Positive value added to the mean square before the reciprocal square root. Defaults to 1e-6.

• bias (bool) – Whether to create an additive bias when affine parameters are enabled. Defaults to False.

• dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Dtype passed to the parameter initializers. Defaults to each initializer’s default dtype.

• rngs (Rngs | None) – Random stream used by parameter initializers. A deterministic stream seeded with zero is used when omitted.

• elementwise_affine (bool) – Whether to create a learnable scale and optional bias with shape num_features. Defaults to True.

• axes (int | Sequence[int] | None) – Input dimensions corresponding to num_features, in the same order. Negative dimensions are supported. Defaults to the trailing feature dimensions.

• scale_initializer (Initializer) – Initializer for the multiplicative scale. Defaults to ones.

• bias_initializer (Initializer) – Initializer for the additive bias. Defaults to zeros.

• quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix configuration. Rules for unrelated operation types are ignored; quantized normalization parameters are not yet supported.

• axis_names (tuple[str | None, ...] | None) – Logical names for all num_features dimensions.

• partition_spec (P | None) – Explicit partition specification for affine parameters. A mapping obtained from axis_names overrides this specification.

• scale_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Metadata attached to the scale parameter.

• bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Metadata attached to the bias parameter.

Variables:

• num_features – Normalized feature shape as a tuple.

• axes – Input dimensions normalized by this module.

• scale – Learnable multiplicative parameter, or None when elementwise_affine=False.

• bias – Learnable additive parameter, or None when disabled.

Example

Normalize the last feature dimension without adding a bias:

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> layer = nn.RMSNorm(4)
>>> x = jnp.arange(8, dtype=jnp.float32).reshape(2, 4)
>>> layer(x).shape
(2, 4)

References

Biao Zhang and Rico Sennrich, “Root Mean Square Layer Normalization” (2019). https://arxiv.org/abs/1910.07467