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:**
  
  - `layer`: The layer to wrap. Usually a [`equinox.nn.Linear`][] or
      a convolutional layer (e.g. [`equinox.nn.Conv2d`][]).
  - `weight_name`: The name of the layer's parameter (a JAX array) to apply
      spectral normalisation to.
  - `num_power_iterations`: The number of power iterations to apply every time
      the array is accessed.
  - `eps`: Epsilon for numerical stability when calculating norms.
  - `inference`: Whether this is in inference mode, at which time no power
      iterations are performed.  This may be toggled with
      [`equinox.nn.inference_mode`][].
  - `key`: A `jax.random.PRNGKey` used to provide randomness for initialisation.
      (Keyword only argument.)
docstr_tok: ["*","*","Arguments",":","*","*","-","`","layer","`",":","The","layer","to","wrap",".","Usually","a","[","`","equinox.nn.Linear","`","]","[","]","or","a","convolutional","layer","(","e.g",".","[","`","equinox.nn.Conv2d","`","]","[","]",")",".","-","`","weight_name","`",":","The","name","of","the","layer","'s","parameter","(","a","JAX","array",")","to","apply","spectral","normalisation","to",".","-","`","num_power_iterations","`",":","The","number","of","power","iterations","to","apply","every","time","the","array","is","accessed",".","-","`","eps","`",":","Epsilon","for","numerical","stability","when","calculating","norms",".","-","`","inference","`",":","Whether","this","is","in","inference","mode",",","at","which","time","no","power","iterations","are","performed",".","This","may","be","toggled","with","[","`","equinox.nn.inference_mode","`","]","[","]",".","-","`","key","`",":","A","`","jax.random.PRNGKey","`","used","to","provide","randomness","for","initialisation",".","(","Keyword","only","argument",".",")"]
code: def __init__(
      self,
      layer: _Layer,
      weight_name: str,
      num_power_iterations: int = 1,
      eps: float = 1e-12,
      inference: bool = False,
      *,
      key: PRNGKeyArray,
  ):
      """**Arguments:**
  
      - `layer`: The layer to wrap. Usually a [`equinox.nn.Linear`][] or
          a convolutional layer (e.g. [`equinox.nn.Conv2d`][]).
      - `weight_name`: The name of the layer's parameter (a JAX array) to apply
          spectral normalisation to.
      - `num_power_iterations`: The number of power iterations to apply every time
          the array is accessed.
      - `eps`: Epsilon for numerical stability when calculating norms.
      - `inference`: Whether this is in inference mode, at which time no power
          iterations are performed.  This may be toggled with
          [`equinox.nn.inference_mode`][].
      - `key`: A `jax.random.PRNGKey` used to provide randomness for initialisation.
          (Keyword only argument.)
      """
  
      self.layer = layer
      self.weight_name = weight_name
      self.num_power_iterations = num_power_iterations
      self.eps = eps
      self.inference = inference
  
      weight = getattr(layer, weight_name)
      if weight.ndim < 2:
          raise ValueError("`weight` must be at least two-dimensional")
      weight = jnp.reshape(weight, (weight.shape[0], -1))
      u_len, v_len = weight.shape
      ukey, vkey = jr.split(key)
      u0 = jr.normal(ukey, (u_len,))
      v0 = jr.normal(vkey, (v_len,))
      for _ in range(15):
          u0, v0 = _power_iteration(weight, u0, v0, eps)
      self.uv_index = StateIndex((u0, v0))
  
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 1 to create a new SpectralNorm object from inherited base classes: StatefulLayer, Generic with example: obj = SpectralNorm(layer, weight_name, num_power_iterations, eps, inference)