keras/src/layers/normalization/batch_normalization.py/BatchNormalization/call
class BatchNormalization:    def call(self, inputs, training=None, mask=None):
        # Check if the mask has one less dimension than the inputs.
        if mask is not None:
            if len(mask.shape) != len(inputs.shape) - 1:
                # Raise a value error
                raise ValueError(
                    "The mask provided should be one dimension less "
                    "than the inputs. Received: "
                    f"mask.shape={mask.shape}, inputs.shape={inputs.shape}"
                )

        input_dtype = standardize_dtype(inputs.dtype)
        if input_dtype in ("float16", "bfloat16"):
            # BN is prone to overflowing for float16/bfloat16 inputs, so we opt
            # out BN for mixed precision.
            inputs = ops.cast(inputs, "float32")

        if training and self.trainable:
            mean, variance = self._moments(inputs, mask)
            moving_mean = ops.cast(self.moving_mean, inputs.dtype)
            moving_variance = ops.cast(self.moving_variance, inputs.dtype)
            self.moving_mean.assign(
                moving_mean * self.momentum + mean * (1.0 - self.momentum)
            )
            self.moving_variance.assign(
                moving_variance * self.momentum
                + variance * (1.0 - self.momentum)
            )
        else:
            moving_mean = ops.cast(self.moving_mean, inputs.dtype)
            moving_variance = ops.cast(self.moving_variance, inputs.dtype)
            mean = moving_mean
            variance = moving_variance

        if self.scale:
            gamma = ops.cast(self.gamma, inputs.dtype)
        else:
            gamma = None

        if self.center:
            beta = ops.cast(self.beta, inputs.dtype)
        else:
            beta = None

        outputs = ops.batch_normalization(
            x=inputs,
            mean=mean,
            variance=variance,
            axis=self.axis,
            offset=beta,
            scale=gamma,
            epsilon=self.epsilon,
        )
        if input_dtype in ("float16", "bfloat16"):
            outputs = ops.cast(outputs, input_dtype)
        return outputs