keras/src/layers/normalization/batch_normalization.py/BatchNormalization/_moments
class BatchNormalization:    def _moments(self, inputs, mask):
        if mask is None:
            return ops.moments(
                inputs,
                axes=self._reduction_axes,
                synchronized=self.synchronized,
            )

        mask_weights = ops.cast(
            mask,
            inputs.dtype,
        )
        mask_weights_broadcasted = ops.expand_dims(
            mask_weights,
            axis=-1,
        )
        weighted_inputs = mask_weights_broadcasted * inputs

        weighted_input_sum = ops.sum(
            weighted_inputs,
            self._reduction_axes,
            keepdims=True,
        )
        sum_of_weights = ops.sum(
            mask_weights_broadcasted,
            self._reduction_axes,
            keepdims=True,
        )
        mean = weighted_input_sum / (sum_of_weights + backend.config.epsilon())

        difference = weighted_inputs - mean
        squared_difference = ops.square(difference)
        weighted_distsq = ops.sum(
            mask_weights_broadcasted * squared_difference,
            self._reduction_axes,
            keepdims=True,
        )
        variance = weighted_distsq / (sum_of_weights + backend.config.epsilon())

        return ops.squeeze(mean), ops.squeeze(variance)