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: Helper to initialize batch norm weights and state.
docstr_tok: ["Helper","to","initialize","batch","norm","weights","and","state","."]
code: def init_weights_and_state(self, input_signature):
      """Helper to initialize batch norm weights and state."""
      axis = self._axis
      axis = (axis,) if jnp.isscalar(axis) else axis
      input_shape = input_signature.shape
      shape = tuple(
          d for i, d in enumerate(input_shape) if i not in axis
      )
      # TODO(jonni): Should beta and gamma match the dtype in the input signature?
      beta = jnp.zeros(shape, dtype="float32") if self._center else ()
      gamma = jnp.ones(shape, dtype="float32") if self._scale else ()
  
      def get_stats_axis(i, d):
          if i in axis:
              return 1
          else:
              return d
  
      stats_shape = tuple(
          get_stats_axis(i, d) for i, d in enumerate(input_shape)
      )
      running_mean = jnp.zeros(stats_shape, dtype=jnp.float32)
      running_var = jnp.ones(stats_shape, dtype=jnp.float32)
      n_batches = jnp.zeros((), dtype=jnp.int64)
      self.weights = (beta, gamma)
      self.state = (running_mean, running_var, n_batches)
  
imports: from trax.fastmath import numpy
  from trax.layers import base
label_desc: Use image node_id 3 for calling the BatchNorm obj's underlying member method code with example usage: obj.init_weights_and_state(input_signature) and returns: int, d