BatchNorm

Applies Batch Normalization to the input.

Abstract Signature:

BatchNorm(input: Tensor, running_mean: Tensor, running_var: Tensor, weight: Tensor, bias: Tensor, training: bool = False, momentum: float = 0.1, eps: float = 1e-5)

PyTorchJAX (Core)NumPyKerasTensorFlowApple MLXFlax NNXPaxML / Praxis

PyTorch

API: torch.batch_norm

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: —

Strategy: Custom / Partial

NumPy

API: —

Strategy: Custom / Partial

Keras

API: keras.ops.batch_normalization

Strategy: Direct Mapping

Official Docs ↗

TensorFlow

API: tf.nn.batch_normalization

Strategy: Direct Mapping

Official Docs ↗

Apple MLX

API: mlx.nn.layers.normalization.BatchNorm

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: nnx.nn.normalization.BatchNorm

Strategy: Direct Mapping

Official Docs ↗

PaxML / Praxis

API: —

Strategy: Custom / Partial