class_docstr: Masking-aware Group Normalization (arxiv.org/abs/1803.08494).
  
  This op is similar to batch normalization, but statistics are shared across
  equally-sized groups of channels and not shared across batch dimension.
  Thus, group normalization does not depend on the batch composition and does
  not require maintaining internal state for storing statistics.
  The user should either specify the total number of channel groups or the
  number of channels per group.
  
  Attributes:
    num_groups: Total number of channel groups. The default value of 32 is
      proposed by the original group normalization paper.
    group_size: The number of channels in a group.
    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.
class_docstr_tok: ["Masking-aware","Group","Normalization","(","arxiv.org\/abs\/1803.08494",")",".","This","op","is","similar","to","batch","normalization",",","but","statistics","are","shared","across","equally-sized","groups","of","channels","and","not","shared","across","batch","dimension",".","Thus",",","group","normalization","does","not","depend","on","the","batch","composition","and","does","not","require","maintaining","internal","state","for","storing","statistics",".","The","user","should","either","specify","the","total","number","of","channel","groups","or","the","number","of","channels","per","group",".","Attributes",":","num_groups",":","Total","number","of","channel","groups",".","The","default","value","of","32","is","proposed","by","the","original","group","normalization","paper",".","group_size",":","The","number","of","channels","in","a","group",".","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","."]
docstr: Applies group normalization to the input (arxiv.org/abs/1803.08494).
  
  Args:
    x: Input of shape N...C, where N is a batch dimension and C is a channels
      dimensions. `...` represents an arbitrary number of extra dimensions
      that are used to accumulate statistics over.
    spatial_shape: Per-input spatial shape of *unpadded* input data with
      dimensions (batch, spatial_dims, channels).
  
  Returns:
    Normalized inputs (the same shape as inputs).
docstr_tok: ["Applies","group","normalization","to","the","input","(","arxiv.org\/abs\/1803.08494",")",".","Args",":","x",":","Input","of","shape","N","...","C",",","where","N","is","a","batch","dimension","and","C","is","a","channels","dimensions",".","`","...","`","represents","an","arbitrary","number","of","extra","dimensions","that","are","used","to","accumulate","statistics","over",".","spatial_shape",":","Per-input","spatial","shape","of","*","unpadded","*","input","data","with","dimensions","(","batch",",","spatial_dims",",","channels",")",".","Returns",":","Normalized","inputs","(","the","same","shape","as","inputs",")","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      spatial_shape: Optional[jnp.ndarray] = None,
  ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]:
      """Applies group normalization to the input (arxiv.org/abs/1803.08494).
  
      Args:
        x: Input of shape N...C, where N is a batch dimension and C is a channels
          dimensions. `...` represents an arbitrary number of extra dimensions
          that are used to accumulate statistics over.
        spatial_shape: Per-input spatial shape of *unpadded* input data with
          dimensions (batch, spatial_dims, channels).
  
      Returns:
        Normalized inputs (the same shape as inputs).
      """
      x = jnp.asarray(x, jnp.float32)
      if (self.num_groups is None and self.group_size is None) or (
          self.num_groups is not None and self.group_size is not None
      ):
          raise ValueError(
              "Either `num_groups` or `group_size` should be "
              "specified, but not both of them."
          )
      num_groups = self.num_groups
  
      channels = x.shape[-1]
      if self.group_size is not None:
          if channels % self.group_size != 0:
              raise ValueError(
                  "Number of channels ({}) is not multiple of the "
                  "group size ({}).".format(channels, self.group_size)
              )
          num_groups = channels // self.group_size
  
      if num_groups <= 0 or channels % num_groups != 0:
          raise ValueError(
              "Number of groups ({}) does not divide the number"
              " of channels ({}).".format(num_groups, channels)
          )
  
      input_shape = x.shape
      group_shape = x.shape[:-1] + (
          num_groups,
          x.shape[-1] // num_groups,
      )
      x = x.reshape(group_shape)
  
      reduction_axis = tuple(range(1, x.ndim - 2)) + (x.ndim - 1,)
      mean = jnp.mean(x, axis=reduction_axis, keepdims=True)
      mean_of_squares = jnp.mean(
          jnp.square(x), axis=reduction_axis, keepdims=True
      )
      orig_denom = np.prod([x.shape[i] for i in reduction_axis[:-1]])
      if (spatial_shape is not None) and self.spatial_norm:
          reduction_axis_shifted = tuple(
              i - 1 for i in reduction_axis[:-1] if i > 0
          )
          denom = jnp.prod(
              spatial_shape[:, reduction_axis_shifted], axis=-1
          )
          denom = jnp.reshape(
              denom, (denom.shape[0],) + (1,) * (mean.ndim - 1)
          )
          denom = jnp.maximum(denom, 1.0)
          mean = mean * (orig_denom / denom)
          mean_of_squares = mean_of_squares * (orig_denom / denom)
  
      var = mean_of_squares - jnp.square(mean)
      x = (x - mean) * lax.rsqrt(var + self.epsilon)
      x = x.reshape(input_shape)
  
      feature_shape = tuple(
          [1 for d in input_shape[:-1]] + [input_shape[-1]]
      )
      if self.use_scale:
          x *= self.param("scale", self.scale_init, feature_shape)
      if self.use_bias:
          x += self.param("bias", self.bias_init, feature_shape)
  
      if spatial_shape is not None:
          # Restore spatial mask for the outputs.
          x = apply_spatial_mask(x, spatial_shape)
      return x.astype(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 GroupNorm obj's underlying member method code with example usage: obj.__call__(x, spatial_shape) and returns: x,  spatial_shape