class_docstr: Layer that performs batch normalization.
  
  In training, batch normalization keeps smoothed cumulative statistics across
  batches of input data and modifies each new batch so that its components are
  normally distributed. In eval or inference, a `BatchNorm` instance uses its
  stored mean and variance to approximately normalize each new batch of data.
  
  See https://arxiv.org/abs/1502.03167 for original presentation and motivation
  of batch normalization).
class_docstr_tok: ["Layer","that","performs","batch","normalization",".","In","training",",","batch","normalization","keeps","smoothed","cumulative","statistics","across","batches","of","input","data","and","modifies","each","new","batch","so","that","its","components","are","normally","distributed",".","In","eval","or","inference",",","a","`","BatchNorm","`","instance","uses","its","stored","mean","and","variance","to","approximately","normalize","each","new","batch","of","data",".","See","https",":","\/\/arxiv.org\/abs\/1502.03167","for","original","presentation","and","motivation","of","batch","normalization",")","."]
docstr: Computes batch normalization as part of a forward pass in the model.
docstr_tok: ["Computes","batch","normalization","as","part","of","a","forward","pass","in","the","model","."]
code: 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
  
imports: from trax.fastmath import numpy
  from trax.layers import base
label_desc: Use image node_id 2 for calling the BatchNorm obj's underlying member method code with example usage: obj.forward(x) and returns: output