class_docstr: Applies spectral normalisation to a given parameter.
  
  Given a weight matrix $W$, and letting $σ(W)$ denote (an approximation to) its
  largest singular value, then this computes $W/σ(W)$.
  
  The approximation $σ(W)$ is computed using
  [power iterations](https://en.wikipedia.org/wiki/Power_iteration)
  which are updated (as a side-effect) every time $W/σ(W)$ is computed.
  
  Spectral normalisation is particularly commonly used when training generative
  adversarial networks; see
  [Spectral Normalization for Generative Adversarial Networks](https://arxiv.org/abs/1802.05957)
  for more details and motivation.
  
  !!! example
  
      See [this example](../../examples/stateful.ipynb) for example usage.
  
  Note that this layer behaves differently during training and inference. During
  training then power iterations are updated; during inference they are fixed.
  Whether the model is in training or inference mode should be toggled using
  [`equinox.nn.inference_mode`][].
docstr: **Arguments:**
  
  - `x`: A JAX array.
  - `state`: An [`equinox.nn.State`][] object (which is used to store the
      power iterations).
  - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
      (Keyword only argument.)
  - `inference`: As per [`equinox.nn.SpectralNorm.__init__`][]. If
      `True` or `False` then it will take priority over `self.inference`. If
      `None` then the value from `self.inference` will be used.
  
  **Returns:**
  
  A 2-tuple of:
  
  - The JAX array from calling `self.layer(x)` (with spectral normalisation
      applied).
  - An updated context object (storing the updated power iterations).
docstr_tok: ["*","*","Arguments",":","*","*","-","`","x","`",":","A","JAX","array",".","-","`","state","`",":","An","[","`","equinox.nn.State","`","]","[","]","object","(","which","is","used","to","store","the","power","iterations",")",".","-","`","key","`",":","Ignored",";","provided","for","compatibility","with","the","rest","of","the","Equinox","API",".","(","Keyword","only","argument",".",")","-","`","inference","`",":","As","per","[","`","equinox.nn.SpectralNorm.__init__","`","]","[","]",".","If","`","True","`","or","`","False","`","then","it","will","take","priority","over","`","self.inference","`",".","If","`","None","`","then","the","value","from","`","self.inference","`","will","be","used",".","*","*","Returns",":","*","*","A","2-tuple","of",":","-","The","JAX","array","from","calling","`","self.layer","(","x",")","`","(","with","spectral","normalisation","applied",")",".","-","An","updated","context","object","(","storing","the","updated","power","iterations",")","."]
code: def __call__(
      self,
      x: Array,
      state: State,
      *,
      key: Optional[PRNGKeyArray] = None,
      inference: Optional[bool] = None,
  ) -> tuple[Array, State]:
      """**Arguments:**
  
      - `x`: A JAX array.
      - `state`: An [`equinox.nn.State`][] object (which is used to store the
          power iterations).
      - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
          (Keyword only argument.)
      - `inference`: As per [`equinox.nn.SpectralNorm.__init__`][]. If
          `True` or `False` then it will take priority over `self.inference`. If
          `None` then the value from `self.inference` will be used.
  
      **Returns:**
  
      A 2-tuple of:
  
      - The JAX array from calling `self.layer(x)` (with spectral normalisation
          applied).
      - An updated context object (storing the updated power iterations).
      """
  
      u, v = state.get(self.uv_index)
      weight = getattr(self.layer, self.weight_name)
      weight_shape = weight.shape
      weight = jnp.reshape(weight, (weight.shape[0], -1))
      if inference is None:
          inference = self.inference
      if not inference:
          stop_weight = lax.stop_gradient(weight)
          for _ in range(self.num_power_iterations):
              u, v = _power_iteration(stop_weight, u, v, self.eps)
          state = state.set(self.uv_index, (u, v))
      σ = jnp.einsum("i,ij,j->", u, weight, v)
      σ_weight = jnp.reshape(weight / σ, weight_shape)
      layer = tree_at(
          lambda l: getattr(l, self.weight_name), self.layer, σ_weight
      )
      out = layer(x)
      return out, state
  
file_path: equinox/equinox/nn/_spectral_norm.py
imports: from typing import Generic, Optional, TypeVar
  import jax
  import jax.lax
  import jax.numpy
  import jax.random
  from jaxtyping import Array, Float, PRNGKeyArray
  from .._module import field
  from .._tree import tree_at
  from ._sequential import StatefulLayer
  from ._stateful import State, StateIndex
label_desc: Use image node_id 2 for calling the SpectralNorm obj's underlying member method code with example usage: obj.__call__(x, state) and returns: out,  state