def forward(self, x):
    """Computes batch normalization as part of a forward pass in the model."""
    running_mean, running_var, n_batches = self.state
    if self._mode == 'train':
      n_batches += 1
      mean, var = self._fast_mean_and_variance(x)
      # Gather smoothed input statistics for later use in evals or inference.
      running_mean = _exponentially_smoothed(self._momentum, running_mean, mean)
      running_var = _exponentially_smoothed(self._momentum, running_var, var)
      self.state = (running_mean, running_var, n_batches)
    else:
      mean = running_mean
      var = running_var

    z = self._z_score(x, mean, var)
    beta, gamma = self._beta_gamma_with_correct_axes(x, self.weights)

    # Return the z rescaled by the parameters if requested.
    if self._center and self._scale:
      output = gamma * z + beta
    elif self._center:
      output = z + beta
    elif self._scale:
      output = gamma * z
    else:
      output = z
    if output.dtype != x.dtype:
      raise TypeError(f'The dtype of the output ({output.dtype}) of batch '
                      f'norm is not the same as the input ({x.dtype}). '
                      f'Batch norm should not change the dtype.')
    return output
Computes batch normalization as part of a forward pass in the model.
trax/layers/normalization.py