class_docstr: A version of BatchNorm that does not track running statistics.
  
  For use in places where this functionality is not available in Jax.
  Attributes:
    axis: the feature or non-batch axis of the input.
    epsilon: a small float added to variance to avoid dividing by zero.
    dtype: the dtype of the computation (default: float32).
    use_bias:  if True, bias (beta) is added.
    use_scale: if True, multiply by scale (gamma). When the next layer is linear
      (also e.g. nn.relu), this can be disabled since the scaling will be done
      by the next layer.
    bias_init: initializer for bias, by default, zero.
    scale_init: initializer for scale, by default, one.
    axis_name: the axis name used to combine batch statistics from multiple
      devices. See `jax.pmap` for a description of axis names (default: None).
    axis_index_groups: groups of axis indices within that named axis
      representing subsets of devices to reduce over (default: None). For
      example, `[[0, 1], [2, 3]]` would independently batch-normalize over the
      examples on the first two and last two devices. See `jax.lax.psum` for
      more details.
docstr: Normalizes the input using batch statistics.
  
  NOTE:
  During initialization (when parameters are mutable) the running average
  of the batch statistics will not be updated. Therefore, the inputs
  fed during initialization don't need to match that of the actual input
  distribution and the reduction axis (set with `axis_name`) does not have
  to exist.
  Args:
    x: the input to be normalized.
    use_running_average: if true, the statistics stored in batch_stats will be
      used instead of computing the batch statistics on the input.
  
  Returns:
    Normalized inputs (the same shape as inputs).
docstr_tok: ["Normalizes","the","input","using","batch","statistics",".","NOTE",":","During","initialization","(","when","parameters","are","mutable",")","the","running","average","of","the","batch","statistics","will","not","be","updated",".","Therefore",",","the","inputs","fed","during","initialization","do","n't","need","to","match","that","of","the","actual","input","distribution","and","the","reduction","axis","(","set","with","`","axis_name","`",")","does","not","have","to","exist",".","Args",":","x",":","the","input","to","be","normalized",".","use_running_average",":","if","true",",","the","statistics","stored","in","batch_stats","will","be","used","instead","of","computing","the","batch","statistics","on","the","input",".","Returns",":","Normalized","inputs","(","the","same","shape","as","inputs",")","."]
code: def __call__(self, x, use_running_average: Optional[bool] = None):
      """Normalizes the input using batch statistics.
  
      NOTE:
      During initialization (when parameters are mutable) the running average
      of the batch statistics will not be updated. Therefore, the inputs
      fed during initialization don't need to match that of the actual input
      distribution and the reduction axis (set with `axis_name`) does not have
      to exist.
      Args:
        x: the input to be normalized.
        use_running_average: if true, the statistics stored in batch_stats will be
          used instead of computing the batch statistics on the input.
  
      Returns:
        Normalized inputs (the same shape as inputs).
      """
      x = jnp.asarray(x, jnp.float32)
      axis = self.axis if isinstance(self.axis, tuple) else (self.axis,)
      axis = _absolute_dims(x.ndim, axis)
      feature_shape = tuple(
          d if i in axis else 1 for i, d in enumerate(x.shape)
      )
      reduced_feature_shape = tuple(
          d for i, d in enumerate(x.shape) if i in axis
      )
      reduction_axis = tuple(i for i in range(x.ndim) if i not in axis)
  
      # see NOTE above on initialization behavior
      initializing = self.is_mutable_collection("params")
  
      mean = jnp.mean(x, axis=reduction_axis, keepdims=False)
      mean2 = jnp.mean(
          lax.square(x), axis=reduction_axis, keepdims=False
      )
      if self.axis_name is not None and not initializing:
          concatenated_mean = jnp.concatenate([mean, mean2])
          mean, mean2 = jnp.split(
              lax.pmean(
                  concatenated_mean,
                  axis_name=self.axis_name,
                  axis_index_groups=self.axis_index_groups,
              ),
              2,
          )
      var = mean2 - lax.square(mean)
  
      y = x - mean.reshape(feature_shape)
      mul = lax.rsqrt(var + self.epsilon)
      if self.use_scale:
          scale = self.param(
              "scale", self.scale_init, reduced_feature_shape
          ).reshape(feature_shape)
          mul = mul * scale
      y = y * mul
      if self.use_bias:
          bias = self.param(
              "bias", self.bias_init, reduced_feature_shape
          ).reshape(feature_shape)
          y = y + bias
      return jnp.asarray(y, self.dtype)
  
imports: import collections
  import functools
  import time
  from typing import Any, Callable, Optional, Tuple
  from dopamine.jax.networks import preprocess_atari_inputs
  from flax import linen
  import gin
  import jax
  from jax import lax
  from jax import random
  from jax.nn import initializers
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the NoStatsBatchNorm obj's underlying member method code with example usage: obj.__call__(x, use_running_average) and returns: jnp