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.

x Input tensor.
mean A mean vector of the same length as the axis dimension of the input thensor.
variance A variance vector of the same length as the axis dimension of the input tensor.
axis Integer, the axis that should be normalized.
offset An 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.
scale A 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.
epsilon Small 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]])