class_docstr: Normalizes an input by its first singular value.
  
  This module uses power iteration to calculate this value based on the
  input and an internal hidden state.
class_docstr_tok: ["Normalizes","an","input","by","its","first","singular","value",".","This","module","uses","power","iteration","to","calculate","this","value","based","on","the","input","and","an","internal","hidden","state","."]
docstr: Performs Spectral Normalization and returns the new value.
  
  Args:
    value: The array-like object for which you would like to perform an
      spectral normalization on.
    update_stats: A boolean defaulting to True. Regardless of this arg, this
      function will return the normalized input. When
      `update_stats` is True, the internal state of this object will also be
      updated to reflect the input value. When `update_stats` is False the
      internal stats will remain unchanged.
    error_on_non_matrix: Spectral normalization is only defined on matrices.
      By default, this module will return scalars unchanged and flatten
      higher-order tensors in their leading dimensions. Setting this flag to
      True will instead throw errors in those cases.
  Returns:
    The input value normalized by it's first singular value.
  Raises:
    ValueError: If `error_on_non_matrix` is True and `value` has ndims > 2.
docstr_tok: ["Performs","Spectral","Normalization","and","returns","the","new","value",".","Args",":","value",":","The","array-like","object","for","which","you","would","like","to","perform","an","spectral","normalization","on",".","update_stats",":","A","boolean","defaulting","to","True",".","Regardless","of","this","arg",",","this","function","will","return","the","normalized","input",".","When","`","update_stats","`","is","True",",","the","internal","state","of","this","object","will","also","be","updated","to","reflect","the","input","value",".","When","`","update_stats","`","is","False","the","internal","stats","will","remain","unchanged",".","error_on_non_matrix",":","Spectral","normalization","is","only","defined","on","matrices",".","By","default",",","this","module","will","return","scalars","unchanged","and","flatten","higher-order","tensors","in","their","leading","dimensions",".","Setting","this","flag","to","True","will","instead","throw","errors","in","those","cases",".","Returns",":","The","input","value","normalized","by","it","'s","first","singular","value",".","Raises",":","ValueError",":","If","`","error_on_non_matrix","`","is","True","and","`","value","`","has","ndims",">","2","."]
code: def __call__(
      self,
      value,
      update_stats: bool = True,
      error_on_non_matrix: bool = False,
  ) -> jax.Array:
      """Performs Spectral Normalization and returns the new value.
  
      Args:
        value: The array-like object for which you would like to perform an
          spectral normalization on.
        update_stats: A boolean defaulting to True. Regardless of this arg, this
          function will return the normalized input. When
          `update_stats` is True, the internal state of this object will also be
          updated to reflect the input value. When `update_stats` is False the
          internal stats will remain unchanged.
        error_on_non_matrix: Spectral normalization is only defined on matrices.
          By default, this module will return scalars unchanged and flatten
          higher-order tensors in their leading dimensions. Setting this flag to
          True will instead throw errors in those cases.
      Returns:
        The input value normalized by it's first singular value.
      Raises:
        ValueError: If `error_on_non_matrix` is True and `value` has ndims > 2.
      """
      value = jnp.asarray(value)
      value_shape = value.shape
  
      # Handle scalars.
      if value.ndim <= 1:
          raise ValueError(
              "Spectral normalization is not well defined for "
              "scalar or vector inputs."
          )
      # Handle higher-order tensors.
      elif value.ndim > 2:
          if error_on_non_matrix:
              raise ValueError(
                  f"Input is {value.ndim}D but error_on_non_matrix is True"
              )
          else:
              value = jnp.reshape(value, [-1, value.shape[-1]])
  
      u0 = hk.get_state(
          "u0",
          [1, value.shape[-1]],
          value.dtype,
          init=hk.initializers.RandomNormal(),
      )
  
      # Power iteration for the weight's singular value.
      for _ in range(self.n_steps):
          v0 = _l2_normalize(
              jnp.matmul(u0, value.transpose([1, 0])), eps=self.eps
          )
          u0 = _l2_normalize(jnp.matmul(v0, value), eps=self.eps)
  
      u0 = jax.lax.stop_gradient(u0)
      v0 = jax.lax.stop_gradient(v0)
  
      sigma = jnp.matmul(jnp.matmul(v0, value), jnp.transpose(u0))[0, 0]
  
      value /= sigma
      value_bar = value.reshape(value_shape)
  
      if update_stats:
          hk.set_state("u0", u0)
          hk.set_state("sigma", sigma)
  
      return value_bar
  
imports: import re
  from typing import Optional
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import initializers
  from haiku._src import module
  import jax
  import jax.lax
  import jax.numpy
label_desc: Use image node_id 2 for calling the SpectralNorm obj's underlying member method code with example usage: obj.__call__(value, update_stats, error_on_non_matrix) and returns: value_bar