class_docstr: Masking-aware Batch Normalization layer.
  
  Attributes:
    use_running_average: If True, the statistics stored in batch_stats
      will be used instead of computing the batch statistics on the input.
    axis: The feature or non-batch axis of the input.
    momentum: Decay rate for the exponential moving average of
      the batch statistics.
    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.
    spatial_norm: If True, spatial shapes influence group norm weights,
      otherwise every batch element has an equal weight.
    axis_name: 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.
class_docstr_tok: ["Masking-aware","Batch","Normalization","layer",".","Attributes",":","use_running_average",":","If","True",",","the","statistics","stored","in","batch_stats","will","be","used","instead","of","computing","the","batch","statistics","on","the","input",".","axis",":","The","feature","or","non-batch","axis","of","the","input",".","momentum",":","Decay","rate","for","the","exponential","moving","average","of","the","batch","statistics",".","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",".","spatial_norm",":","If","True",",","spatial","shapes","influence","group","norm","weights",",","otherwise","every","batch","element","has","an","equal","weight",".","axis_name",":","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.
  
  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.
    spatial_shape: Per-input spatial shape of *unpadded* input data with
      dimensions (batch, spatial_dims, channels).
  
  Returns:
    Normalized inputs (the same shape as inputs) and the spatial shape.
docstr_tok: ["Normalizes","the","input","using","batch","statistics",".","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",".","spatial_shape",":","Per-input","spatial","shape","of","*","unpadded","*","input","data","with","dimensions","(","batch",",","spatial_dims",",","channels",")",".","Returns",":","Normalized","inputs","(","the","same","shape","as","inputs",")","and","the","spatial","shape","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      use_running_average: Optional[bool] = None,
      spatial_shape: Optional[jnp.ndarray] = None,
  ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]:
      """Normalizes the input using batch statistics.
  
      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.
        spatial_shape: Per-input spatial shape of *unpadded* input data with
          dimensions (batch, spatial_dims, channels).
  
      Returns:
        Normalized inputs (the same shape as inputs) and the spatial shape.
      """
      use_running_average = nn.module.merge_param(
          "use_running_average",
          self.use_running_average,
          use_running_average,
      )
      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)
  
      # Detect if we're in initialization via empty variable tree.
      initializing = not self.has_variable("batch_stats", "mean")
  
      ra_mean = self.variable(
          "batch_stats",
          "mean",
          lambda s: jnp.zeros(s, jnp.float32),
          reduced_feature_shape,
      )
      ra_var = self.variable(
          "batch_stats",
          "var",
          lambda s: jnp.ones(s, jnp.float32),
          reduced_feature_shape,
      )
  
      if use_running_average:
          mean, var = ra_mean.value, ra_var.value
      else:
          p_agg = (
              (self.axis_name is not None)
              and (not initializing)
              and self.spatial_norm
          )
          mean, var = _bn_agg_mean_var(
              x,
              reduction_axis,
              p_agg,
              axis_name=self.axis_name,
              axis_index_groups=self.axis_index_groups,
              spatial_shape=spatial_shape,
          )
  
          if not initializing:
              ra_mean.value = (
                  self.momentum * ra_mean.value
                  + (1 - self.momentum) * mean
              )
              ra_var.value = (
                  self.momentum * ra_var.value
                  + (1 - self.momentum) * var
              )
  
      # Apply normaliation.
      mean, var = mean.reshape(feature_shape), var.reshape(
          feature_shape
      )
      y = x - mean
      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
  
      if spatial_shape is not None:
          # Restore spatial mask for the outputs.
          y = apply_spatial_mask(y, spatial_shape)
      return jnp.asarray(y, self.dtype), spatial_shape
  
imports: import functools
  from typing import Optional, Tuple, Union, Callable, Any, Sequence, List, Iterable
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the BatchNorm obj's underlying member method code with example usage: obj.__call__(x, use_running_average, spatial_shape) and returns: jnp,  spatial_shape