tf.keras.ops.batch_normalization

Normalizes x by mean and variance.

tf.keras.ops.batch_normalization(
    x, mean, variance, axis, offset=None, scale=None, epsilon=0.001
)

This op is typically used by the batch normalization step in a neural network. It normalizes the input tensor along the given axis.

xInput tensor.
meanA mean vector of the same length as the axis dimension of the input thensor.
varianceA variance vector of the same length as the axis dimension of the input tensor.
axisInteger, the axis that should be normalized.
offsetAn offset vector of the same length as the axis dimension of the input tensor. If not None, offset is added to the normalized tensor. Defaults to None.
scaleA scale vector of the same length as the axis dimension of the input tensor. If not None, the normalized tensor is multiplied by scale. Defaults to None.
epsilonSmall float added to variance to avoid dividing by zero. Defaults to 1e-3.

The normalized tensor.

Example:

x = keras.ops.convert_to_tensor(
    [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
)
keras.ops.batch_normalization(
    x,
    mean=[0.4, 0.5, 0.6],
    variance=[0.67, 0.67, 0.67],
    axis=-1
)
array([[-3.6624e-01, -3.6624e-01, -3.6624e-01],
       [-4.6445e-09,  0.0000e+00, -1.8578e-08],
       [ 3.6624e-01,  3.6624e-01,  3.6624e-01]])