Normalizer Reference

Data normalization is essential for the stable and efficient training of Neural Operators. Skewed training data distributions heavily influence model performance and training, especially for physical data which may range across vastly different scales. neojax provides multiple normalizers for reliable training.

All normalizers share a common API, inherit from BaseNormalizer, and can be arbitrarily composed to build complex normalization pipelines.

The Normalizer API

We adopt the common naming conventions of scientific computing libraries. Every normalizer implements:

• compute_stats(data, axis=...): 'Learns' the normalization statistics (e.g., min/max, mean/std) from the given data. Because neojax uses equinox and is strictly functional, this returns a new normalizer instance with the updated statistics.
• transform(x) / __call__(x): Applies the normalization to the input.
• inverse_transform(x): Reverses the normalization, projecting predictions back into the original physical space.

Example Usage

import jax.numpy as jnp
from neojax.data.normalizers import MinMaxNormalizer

# 1. Initialize the normalizer
normalizer = MinMaxNormalizer()

# 2. Learn statistics from training data
train_data = jnp.array([[-10.0, 0.0, 10.0]])
normalizer = normalizer.compute_stats(train_data)

# 3. Transform data before passing to the model
normalized_data = normalizer(train_data) 
# normalized_data is now scaled to [0, 1]

# 4. Revert model predictions back to the original scale
predictions = model(normalized_data)
physical_predictions = normalizer.inverse_transform(predictions)

Composing Normalizers

Data pipelines often require multiple transformation steps. The ComposedNormalizer allows you to pipe normalizers together.

By default, compute_stats works sequentially through the pipeline: the second normalizer computes its statistics on the data after it has been transformed by the first normalizer. You can disable this by passing sequential=False.

from neojax.data.normalizers import ComposedNormalizer, RobustNormalizer, MinMaxNormalizer

# Build a pipeline
pipeline = ComposedNormalizer(
    RobustNormalizer(),
    MinMaxNormalizer()
)

# Computes robust stats, transforms data, then computes min/max stats on the output
pipeline = pipeline.compute_stats(data)

Base Normalizer

neojax.data.normalizers.base_normalizer.BaseNormalizer

Base class for all normalizers.

All normalizers should inherit from BaseNormalizer. Following the "Abstract or Final" pattern, this class contains no logic and only defines the interface. Subclasses should be marked as final.

Normalizer methods follow the naming conventions of major libraries, where transform is the forward operation and inverse_transform the inverse.

Internal Attributes

These fields store the internal state of the normalizer.

• stats (PyTree): A PyTree containing the normalization statistics (e.g. mean, std).

compute_stats(data: Float[Array, ...]) -> BaseNormalizer

Computes statistics from the provided data.

This function should set the stats attribute.

Parameters:

Name	Type	Description	Default
data	Float[Array, ...]	The data to compute statistics from.	required

Returns:

Type	Description
BaseNormalizer	A new BaseNormalizer instance with updated stats.

transform(x: Float[Array, ...]) -> Float[Array, ...]

Applies the forward normalization transform.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The input array to transform.	required

Returns:

Type	Description
Float[Array, ...]	The transformed array.

inverse_transform(x: Float[Array, ...]) -> Float[Array, ...]

Applies the inverse normalization transform.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The transformed array to revert.	required

Returns:

Type	Description
Float[Array, ...]	The array in its original scale

Unit Gaussian Normalizer

neojax.data.normalizers.unit_gaussian_normalizer.UnitGaussianNormalizer

Normalizes data to be mean zero and unit std.

Parameters:

Name	Type	Description	Default
mean	float | Float[Array, ...]	Optional initial mean if known. Default is 0.0.	0.0
std	float | Float[Array, ...]	Optional initial standard deviation if known. Default is unit 1.0.	1.0

Internal Attributes

These fields store the internal state of the normalizer.

• stats (dict[str, Float[Array, ...] | None]): Dict of mean and std.

mean property

Accessor for the mean stored in stats.

std property

Accessor for the std stored in stats.

compute_stats(data: Float[Array, 'c ...'], axis: int | tuple[int, ...] | None = None) -> UnitGaussianNormalizer

Computes mean and std from diven data.

To normalize per-channel for input (c, d1, ..., dN), pass axis=tuple(range(1, data.ndim)). To normalize across all axes, pass axis=None.

Parameters:

Name	Type	Description	Default
data	Float[Array, 'c ...']	The input array.	required
axis	int | tuple[int, ...] | None	Axis or axes along which the statistics are computed. The default (None) computes the global mean/std.	None

Returns:

Type	Description
UnitGaussianNormalizer	A new UnitGaussianNormalizer instance
UnitGaussianNormalizer	with updated mean and std.

transform(x: Float[Array, ...]) -> Float[Array, ...]

Standardizes input using stored mean and std.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The input array to transform.	required

Returns:

Type	Description
Float[Array, ...]	The standardized array.

inverse_transform(x: Float[Array, ...]) -> Float[Array, ...]

Reverts standardization.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The standardized array to revert.	required

Returns:

Type	Description
Float[Array, ...]	The de-standardized array.

Robust Normalizer

neojax.data.normalizers.robust_normalizer.RobustNormalizer

Normalizes data using median and interquartile range.

This normalizer is robust to outliers as it uses the median and the interquartile range (IQR).

Parameters:

Name	Type	Description	Default
median	float | Float[Array, ...] | None	Optional median for dynamical scaling. Default is None, i.e., use learned median from self.stats.	None
scale	float | Float[Array, ...] | None	Optional interquartile range for dynamical scaling. Default is None, i.e., use learned iqr from self.stats.	None
quantile_range	tuple[float, float]		(25.0, 75.0)

Internal Attributes

These fields store the internal state of the normalizer.

• quantile_range (tuple[float, float]): The lower and upper quantile bounds.
• stats (dict[str, Float[Array, ...] | None]): Dict containing 'median' and 'scale' (iqr).

median property

Accessor for the median stored in stats.

scale property

Accessor for the scale stored in stats.

compute_stats(data: Float[Array, 'c ...'], axis: int | tuple[int, ...] | None = None) -> RobustNormalizer

Computes median and quantile range from given data.

Parameters:

Name	Type	Description	Default
data	Float[Array, 'c ...']	The input array to compute statistics from.	required
axis	int | tuple[int, ...] | None	Axis or axes along which the statistics are computed. The default (None) computes the global statistics.	None

Returns:

Type	Description
RobustNormalizer	A new RobustNormalizer instance with updated statistics.

transform(x: Float[Array, ...]) -> Float[Array, ...]

Normalizes input using median and quantile range.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The input array to transform.	required

Returns:

Type	Description
Float[Array, ...]	The normalized array.

inverse_transform(x: Float[Array, ...]) -> Float[Array, ...]

Reverts robust normalization.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The transformed array to revert.	required

Returns:

Type	Description
Float[Array, ...]	The de-normalized array.

Min/Max Normalizer

(Supports both "scale" and destructive "clip" modes)

neojax.data.normalizers.min_max_normalizer.MinMaxNormalizer

Normalizes data to be within a given min and max.

Minima and maxima can be passed the following ways:

1. During initialization if known or explicitly required (e.g. known physical bounds).
2. Learned from data by passing them to compute_stats

Data can be clipped or scaled to [min, max] range.

Parameters:

Name	Type	Description	Default
minima	float | int | Float[Array, ...] | None	Optional min for dynamical scaling. Default is None, i.e., use learned min from self.stats. If Array, should be broadcast-compatible with transformation data and maxima.	None
maxima	float | int | Float[Array, ...] | None	Optional max for dynamical scaling. Default is None, i.e., use learned max from self.stats. If Array, should be broadcast-compatible with transformation data and minima.	None

Internal Attributes

These fields store the internal state of the normalizer.

• stats (dict[str, Float[Array, ...] | None]): Dict of min and max if passed at initialization, else None.

Warning

Inverse transformation can only be applied for scaled transformations.

maxima property

Accessor for the max stored in stats.

minima property

Accessor for the min stored in stats.

compute_stats(data: Float[Array, 'c ...'], axis: int | tuple[int, ...] | None = None) -> MinMaxNormalizer

Computes min and max from given data.

To normalize per-channel for input (c, d1, ..., dN), pass axis=tuple(range(1, data.ndim)). To normalize across all axes, pass axis=None.

Parameters:

Name	Type	Description	Default
data	Float[Array, 'c ...']	The input array.	required
axis	int | tuple[int, ...] | None	Axis or axes along which the statistics are computed. The default (None) computes the global min/max.	None

Returns:

Type	Description
MinMaxNormalizer	A new MinMaxNormalizer instance
MinMaxNormalizer	with updated min and max.

transform(x: Float[Array, ...], mode: typing.Literal['clip', 'scale'] = 'scale') -> Float[Array, ...]

Normalizes input by clipping or rescaling it.

Uses learned minima or maxima from self.stats.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The input array to transform.	required
mode	Literal['clip', 'scale']	Whether to "clip" values to be in [minima, maxima] (inclusive) or "scale" values inside range with standard min/max scaling (x - min) / (max - min + eps). Default is "scale".	'scale'

Note

"clip" is a destructive operation. Clipped data cannot be perfectly reconstructed:

    inverse_transform(transform(x, mode="clip", ...), ...) != x

Returns:

Type	Description
Float[Array, ...]	The min/max normalized array.

Raises:

Type	Description
ValueError	If the transformation mode is unsupported.

inverse_transform(x: Float[Array, ...]) -> Float[Array, ...]

Reverts scaled min/max normalization.

Can only reverse scaled normalization. Clipped normalizations are irreversible.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The standardized array to revert.	required

Returns:

Type	Description
Float[Array, ...]	The de-normalized array.

Physics Normalizer

(Non-dimensionalizes data using physical scales. See Scales Reference.)

neojax.data.normalizers.physics_normalizer.PhysicsNormalizer

Normalizes data based on spatial dimensions/physical units.

This normalizer applies a product of multiple PhysicalScale factors to non-dimensionalize the input data.

Parameters:

Name	Type	Description	Default
*scales	PhysicalScale	Variable number of PhysicalScale providers to be applied to the data.	required

Internal Attributes

These fields store the internal state of the normalizer.

• scales (tuple[PhysicalScale, ...]): Tuple of physical scale providers.
• stats (dict[str, Float[Array, ...] | None]): Dict containing 'scale_product', which is the resulting element-wise product of all computed scales.

Examples:

from neojax.data.scales import CharacteristicLengthScale, ReynoldsScale

# Compose multiple physical scales
norm = PhysicsNormalizer(
    CharacteristicLengthScale(L_ref=1.0),
    ReynoldsScale(U=10.0, L=1.0, nu=1e-3)
)

# Compute and apply
norm = norm.compute_stats(data)
non_dimensional = norm(data)

compute_stats(data: Float[Array, 'c ...'], axis: int | tuple[int, ...] | None = None) -> PhysicsNormalizer

Computes the product of all scales for the given data.

Parameters:

Name	Type	Description	Default
data	Float[Array, 'c ...']	The input array to compute statistics from.	required
axis	int | tuple[int, ...] | None	Unused for physics-based scaling but kept for API compatibility.	None

Returns:

Type	Description
PhysicsNormalizer	A new PhysicsNormalizer instance with the final computed
PhysicsNormalizer	scale product in stats.

transform(x: Float[Array, ...]) -> Float[Array, ...]

Non-dimensionalizes input by dividing by the scale product.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The input array to transform.	required

Returns:

Type	Description
Float[Array, ...]	The non-dimensionalized array.

inverse_transform(x: Float[Array, ...]) -> Float[Array, ...]

Reverts non-dimensionalization by multiplying by the scale product.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The transformed array to revert.	required

Returns:

Type	Description
Float[Array, ...]	The array in physical units.

Composed Normalizer

neojax.data.normalizers.composed_normalizer.ComposedNormalizer

Composes multiple normalizers into a pipeline.

The normalizers are applied sequentially in the order they are provided. This class allows for complex normalization pipelines, such as first applying a robust scaler and then a min-max scaling.

Parameters:

Name	Type	Description	Default
*normalizers	BaseNormalizer	Normalizer instances to be composed in the pipeline.	required

Internal Attributes

These fields store the internal state of the normalizer.

• stats (dict[str, tuple[BaseNormalizer, ...]]): Dict containing 'normalizers', which stores the sequence of normalizers as a tuple.

Examples:

import jax.numpy as jnp
from neojax.data.normalizers import (
    ComposedNormalizer,
    MinMaxNormalizer,
    UnitGaussianNormalizer
)

norm = ComposedNormalizer(
    MinMaxNormalizer(),
    UnitGaussianNormalizer()
)

data = jnp.ones((3, 3, 3))
# compute stats (sequentially)
norm = norm.compute_stats(data)

transformed = norm(data)

Notes

If all normalizer statistics need to be computed on raw data, pass sequential=False to compute_stats().

compute_stats(data: Float[Array, 'c ...'], axis: int | tuple[int, ...] | None = None, sequential: bool = True) -> ComposedNormalizer

Computes statistics for each normalizer in the pipeline.

Each normalizer in the pipeline computes its statistics independently based on the provided data, either sequentially of non-sequentially. See sequential for details.

Parameters:

Name	Type	Description	Default
data	Float[Array, 'c ...']	The input array to compute statistics from.	required
axis	int | tuple[int, ...] | None	Axis or axes along which the statistics are computed. The default (None) computes the global statistics.	None
sequential	bool	Whether to apply stats computations in sequence, i.e., compute the stats on the transformed inputs of the previous normalizer. Non-sequential computes each stat on the raw inputs (e.g. for parallel computation). Default is sequential.	True

Returns:

Type	Description
ComposedNormalizer	A new ComposedNormalizer instance with updated statistics
ComposedNormalizer	for all sub-normalizers.

transform(x: Float[Array, ...]) -> Float[Array, ...]

Applies the sequence of normalizations to the input.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The input array to transform.	required

Returns:

Type	Description
Float[Array, ...]	The transformed array.

inverse_transform(x: Float[Array, ...]) -> Float[Array, ...]

Reverts the sequence of normalizations in reverse order.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The transformed array to revert.	required

Returns:

Type	Description
Float[Array, ...]	The de-normalized array.