class_docstr: Applies weight normalisation to a given parameter.
  
  Given the 2D weight matrix
  $W = (W_{ij})_{ij} \in \mathbb{R}^{\text{out}} \times \mathbb{R}^{\text{in}}$ of a
  linear layer, then it replaces it with the following reparameterisation:
  
  $g \frac{v_{ij}}{\lVert v_{i\, \cdot} \rVert} \in \mathbb{R}^{\text{out}} \times \mathbb{R}^{\text{in}}$
  
  where $v_{ij}$ is initialised as $W_{ij}$, and $g$ is initialised as
  $\lVert v_{i\, \cdot} \rVert = \sum_j {v_{ij}}^2$.
  
  Overall, the direction ($v$) and the magnitude ($g$) of the output of each neuron
  are treated separately.
  
  Given n-dimensional weight matrices $W$ (in convolutional layers), then the
  normalisation is analogusly instead computed over every axis except the first.
  
  ??? cite
  
      [Weight Normalisation](https://arxiv.org/abs/1602.07868)
  
      ```bibtex
      @article{DBLP:journals/corr/SalimansK16,
      author       = {Tim Salimans and
                      Diederik P. Kingma},
      title        = {Weight Normalisation: {A} Simple
                      Reparameterization to Accelerate
                      Training of Deep Neural Networks},
      journal      = {CoRR},
      volume       = {abs/1602.07868},
      year         = {2016},
      url          = {http://arxiv.org/abs/1602.07868},
      eprinttype   = {arXiv},
      eprint       = {1602.07868},
      timestamp    = {Mon, 13 Aug 2018 16:47:07 +0200},
      biburl       = {https://dblp.org/rec/journals/corr/SalimansK16.bib},
      bibsource    = {dblp computer science bibliography, https://dblp.org}
      }
      ```
docstr: **Arguments:**
  
  - `x`: A JAX Array.
  - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
  
  **Returns:**
  
  - The JAX array from calling `self.layer(x)` (with weight normalisation
  applied).
docstr_tok: ["*","*","Arguments",":","*","*","-","`","x","`",":","A","JAX","Array",".","-","`","key","`",":","Ignored",";","provided","for","compatibility","with","the","rest","of","the","Equinox","API",".","*","*","Returns",":","*","*","-","The","JAX","array","from","calling","`","self.layer","(","x",")","`","(","with","weight","normalisation","applied",")","."]
code: def __call__(
      self, x: Array, *, key: Optional[PRNGKeyArray] = None
  ) -> Array:
      """**Arguments:**
  
      - `x`: A JAX Array.
      - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
  
      **Returns:**
  
      - The JAX array from calling `self.layer(x)` (with weight normalisation
      applied).
      """
  
      v = getattr(self.layer, self.weight_name)
      weight = v * self.g / self._norm(v)
      layer = tree_at(
          lambda l: getattr(l, self.weight_name), self.layer, weight
      )
      return layer(x)
  
file_path: equinox/equinox/nn/_weight_norm.py
imports: import functools
  from collections.abc import Callable
  from typing import Generic, Optional, TypeVar
  import jax
  import jax.numpy
  from jaxtyping import Array, PRNGKeyArray, Scalar
  from .._module import field, Module
  from .._tree import tree_at
label_desc: Use image node_id 2 for calling the WeightNorm obj's underlying member method code with example usage: obj.__call__(x) and returns: layer