class_docstr: Spectral normalization. See:
  
  - https://arxiv.org/abs/1802.05957
  - https://arxiv.org/abs/1805.08318
  - https://arxiv.org/abs/1809.11096
  
  Spectral normalization normalizes the weight params so that the spectral
  norm of the matrix is equal to 1. This is implemented as a layer wrapper
  where each wrapped layer will have its params spectral normalized before
  computing its ``__call__`` output.
  
  Usage Note:
  The initialized variables dict will contain, in addition to a 'params'
  collection, a separate 'batch_stats' collection that will contain a
  ``u`` vector and ``sigma`` value, which are intermediate values used
  when performing spectral normalization. During training, we pass in
  ``update_stats=True`` and ``mutable=['batch_stats']`` so that ``u``
  and ``sigma`` are updated with the most recently computed values using
  power iteration. This will help the power iteration method approximate
  the true singular value more accurately over time. During eval, we pass
  in ``update_stats=False`` to ensure we get deterministic behavior from
  the model. For example::
  
  Example usage::
  
    >>> import flax, flax.linen as nn
    >>> import jax, jax.numpy as jnp
    >>> import optax
  
    >>> class Foo(nn.Module):
    ...   @nn.compact
    ...   def __call__(self, x, train):
    ...     x = nn.Dense(3)(x)
    ...     # only spectral normalize the params of the second Dense layer
    ...     x = nn.SpectralNorm(nn.Dense(4))(x, update_stats=train)
    ...     x = nn.Dense(5)(x)
    ...     return x
  
    >>> # init
    >>> x = jnp.ones((1, 2))
    >>> y = jnp.ones((1, 5))
    >>> model = Foo()
    >>> variables = model.init(jax.random.PRNGKey(0), x, train=False)
    >>> flax.core.freeze(jax.tree_map(jnp.shape, variables))
    FrozenDict({
        batch_stats: {
            SpectralNorm_0: {
                Dense_1/kernel/sigma: (),
                Dense_1/kernel/u: (1, 4),
            },
        },
        params: {
            Dense_0: {
                bias: (3,),
                kernel: (2, 3),
            },
            Dense_1: {
                bias: (4,),
                kernel: (3, 4),
            },
            Dense_2: {
                bias: (5,),
                kernel: (4, 5),
            },
        },
    })
  
    >>> # train
    >>> def train_step(variables, x, y):
    ...   def loss_fn(params):
    ...     logits, updates = model.apply(
    ...         {'params': params, 'batch_stats': variables['batch_stats']},
    ...         x,
    ...         train=True,
    ...         mutable=['batch_stats'],
    ...     )
    ...     loss = jnp.mean(optax.l2_loss(predictions=logits, targets=y))
    ...     return loss, updates
    ...
    ...   (loss, updates), grads = jax.value_and_grad(loss_fn, has_aux=True)(
    ...       variables['params']
    ...   )
    ...   return {
    ...       'params': jax.tree_map(
    ...           lambda p, g: p - 0.1 * g, variables['params'], grads
    ...       ),
    ...       'batch_stats': updates['batch_stats'],
    ...   }, loss
    >>> for _ in range(10):
    ...   variables, loss = train_step(variables, x, y)
  
    >>> # inference / eval
    >>> out = model.apply(variables, x, train=False)
  
  Attributes:
    layer_instance: Module instance that is wrapped with SpectralNorm
    n_steps: How many steps of power iteration to perform to approximate the
      singular value of the weight params.
    epsilon: A small float added to l2-normalization to avoid dividing by zero.
    dtype: the dtype of the result (default: infer from input and params).
    param_dtype: the dtype passed to parameter initializers (default: float32).
    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 an error if a weight tensor with dimension
      greater than 2 is used by the layer.
    collection_name: Name of the collection to store intermediate values used
      when performing spectral normalization.
class_docstr_tok: ["Spectral","normalization",".","See",":","-","https",":","\/\/arxiv.org\/abs\/1802.05957","-","https",":","\/\/arxiv.org\/abs\/1805.08318","-","https",":","\/\/arxiv.org\/abs\/1809.11096","Spectral","normalization","normalizes","the","weight","params","so","that","the","spectral","norm","of","the","matrix","is","equal","to","1",".","This","is","implemented","as","a","layer","wrapper","where","each","wrapped","layer","will","have","its","params","spectral","normalized","before","computing","its","``","__call__","``","output",".","Usage","Note",":","The","initialized","variables","dict","will","contain",",","in","addition","to","a","'params'","collection",",","a","separate","'batch_stats","'","collection","that","will","contain","a","``","u","``","vector","and","``","sigma","``","value",",","which","are","intermediate","values","used","when","performing","spectral","normalization",".","During","training",",","we","pass","in","``","update_stats=True","``","and","``","mutable=","[","'batch_stats","'","]","``","so","that","``","u","``","and","``","sigma","``","are","updated","with","the","most","recently","computed","values","using","power","iteration",".","This","will","help","the","power","iteration","method","approximate","the","true","singular","value","more","accurately","over","time",".","During","eval",",","we","pass","in","``","update_stats=False","``","to","ensure","we","get","deterministic","behavior","from","the","model",".","For","example",":",":","Example","usage",":",":",">",">",">","import","flax",",","flax.linen","as","nn",">",">",">","import","jax",",","jax.numpy","as","jnp",">",">",">","import","optax",">",">",">","class","Foo","(","nn.Module",")",":","...","@","nn.compact","...","def","__call__","(","self",",","x",",","train",")",":","...","x","=","nn.Dense","(","3",")","(","x",")","...","#","only","spectral","normalize","the","params","of","the","second","Dense","layer","...","x","=","nn.SpectralNorm","(","nn.Dense","(","4",")",")","(","x",",","update_stats=train",")","...","x","=","nn.Dense","(","5",")","(","x",")","...","return","x",">",">",">","#","init",">",">",">","x","=","jnp.ones","(","(","1",",","2",")",")",">",">",">","y","=","jnp.ones","(","(","1",",","5",")",")",">",">",">","model","=","Foo","(",")",">",">",">","variables","=","model.init","(","jax.random.PRNGKey","(","0",")",",","x",",","train=False",")",">",">",">","flax.core.freeze","(","jax.tree_map","(","jnp.shape",",","variables",")",")","FrozenDict","(","{","batch_stats",":","{","SpectralNorm_0",":","{","Dense_1\/kernel\/sigma",":","(",")",",","Dense_1\/kernel\/u",":","(","1",",","4",")",",","}",",","}",",","params",":","{","Dense_0",":","{","bias",":","(","3",",",")",",","kernel",":","(","2",",","3",")",",","}",",","Dense_1",":","{","bias",":","(","4",",",")",",","kernel",":","(","3",",","4",")",",","}",",","Dense_2",":","{","bias",":","(","5",",",")",",","kernel",":","(","4",",","5",")",",","}",",","}",",","}",")",">",">",">","#","train",">",">",">","def","train_step","(","variables",",","x",",","y",")",":","...","def","loss_fn","(","params",")",":","...","logits",",","updates","=","model.apply","(","...","{","'params","'",":","params",",","'batch_stats","'",":","variables","[","'batch_stats","'","]","}",",","...","x",",","...","train=True",",","...","mutable=","[","'batch_stats","'","]",",","...",")","...","loss","=","jnp.mean","(","optax.l2_loss","(","predictions=logits",",","targets=y",")",")","...","return","loss",",","updates","...","...","(","loss",",","updates",")",",","grads","=","jax.value_and_grad","(","loss_fn",",","has_aux=True",")","(","...","variables","[","'params","'","]","...",")","...","return","{","...","'params","'",":","jax.tree_map","(","...","lambda","p",",","g",":","p","-","0.1","*","g",",","variables","[","'params","'","]",",","grads","...",")",",","...","'batch_stats","'",":","updates","[","'batch_stats","'","]",",","...","}",",","loss",">",">",">","for","_","in","range","(","10",")",":","...","variables",",","loss","=","train_step","(","variables",",","x",",","y",")",">",">",">","#","inference","\/","eval",">",">",">","out","=","model.apply","(","variables",",","x",",","train=False",")","Attributes",":","layer_instance",":","Module","instance","that","is","wrapped","with","SpectralNorm","n_steps",":","How","many","steps","of","power","iteration","to","perform","to","approximate","the","singular","value","of","the","weight","params",".","epsilon",":","A","small","float","added","to","l2-normalization","to","avoid","dividing","by","zero",".","dtype",":","the","dtype","of","the","result","(","default",":","infer","from","input","and","params",")",".","param_dtype",":","the","dtype","passed","to","parameter","initializers","(","default",":","float32",")",".","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","an","error","if","a","weight","tensor","with","dimension","greater","than","2","is","used","by","the","layer",".","collection_name",":","Name","of","the","collection","to","store","intermediate","values","used","when","performing","spectral","normalization","."]
docstr: Compute the largest singular value using power iteration and normalize
  the variables ``vs`` using this value. This is intended to be a helper
  function used in this Module's ``__call__`` method in conjunction with
  ``nn.transforms.map_variables`` and ``jax.tree_util.tree_map_with_path``.
  
  Args:
    path: dict key path, used for naming the ``u`` and ``sigma`` variables
    vs: variables to be spectral normalized
    update_stats: if True, update the ``u`` vector and ``sigma`` variables
      after computing their updated values using power iteration. This will
      help the power iteration method approximate the true singular value
      more accurately over time.
docstr_tok: ["Compute","the","largest","singular","value","using","power","iteration","and","normalize","the","variables","``","vs","``","using","this","value",".","This","is","intended","to","be","a","helper","function","used","in","this","Module","'s","``","__call__","``","method","in","conjunction","with","``","nn.transforms.map_variables","``","and","``","jax.tree_util.tree_map_with_path","``",".","Args",":","path",":","dict","key","path",",","used","for","naming","the","``","u","``","and","``","sigma","``","variables","vs",":","variables","to","be","spectral","normalized","update_stats",":","if","True",",","update","the","``","u","``","vector","and","``","sigma","``","variables","after","computing","their","updated","values","using","power","iteration",".","This","will","help","the","power","iteration","method","approximate","the","true","singular","value","more","accurately","over","time","."]
code: def _spectral_normalize(self, path, vs, update_stats):
      """Compute the largest singular value using power iteration and normalize
      the variables ``vs`` using this value. This is intended to be a helper
      function used in this Module's ``__call__`` method in conjunction with
      ``nn.transforms.map_variables`` and ``jax.tree_util.tree_map_with_path``.
  
      Args:
        path: dict key path, used for naming the ``u`` and ``sigma`` variables
        vs: variables to be spectral normalized
        update_stats: if True, update the ``u`` vector and ``sigma`` variables
          after computing their updated values using power iteration. This will
          help the power iteration method approximate the true singular value
          more accurately over time.
      """
      value = jnp.asarray(vs)
      value_shape = value.shape
  
      # Skip and return value if input is scalar, vector or if number of power
      # iterations is less than 1
      if value.ndim <= 1 or self.n_steps < 1:
          return value
      # Handle higher-order tensors.
      elif value.ndim > 2:
          if self.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]))
  
      u_var_name = (
          self.layer_instance.name
          + "/"
          + "/".join((dict_key.key for dict_key in path[1:]))
          + "/u"
      )
      u_var = self.variable(
          self.collection_name,
          u_var_name,
          jax.random.normal,
          self.make_rng("params")
          if not self.has_variable(self.collection_name, u_var_name)
          else None,
          (1, value.shape[-1]),
          self.param_dtype,
      )
      u0 = u_var.value
      sigma_var_name = (
          self.layer_instance.name
          + "/"
          + "/".join((dict_key.key for dict_key in path[1:]))
          + "/sigma"
      )
      sigma_var = self.variable(
          self.collection_name,
          sigma_var_name,
          jnp.ones,
          (),
          self.param_dtype,
      )
  
      # 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.epsilon
          )
          u0 = _l2_normalize(jnp.matmul(v0, value), eps=self.epsilon)
  
      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 /= jnp.where(sigma != 0, sigma, 1)
      value_bar = value.reshape(value_shape)
  
      if update_stats:
          u_var.value = u0
          sigma_var.value = sigma
  
      dtype = dtypes.canonicalize_dtype(
          vs, u0, v0, sigma, dtype=self.dtype
      )
      return jnp.asarray(value_bar, dtype)
  
imports: import dataclasses
  import functools
  from typing import Any, Callable, Iterable, Optional, Sequence, Tuple, Union
  import jax
  import jax.numpy
  from jax import lax
  from jax.nn import initializers
  from flax.linen import dtypes, module, transforms
label_desc: Use image node_id 2 for calling the SpectralNorm obj's underlying member method code with example usage: obj._spectral_normalize(path, vs, update_stats) and returns: jnp, value