class_docstr: L2 weight normalization (https://arxiv.org/pdf/1602.07868.pdf).
  
  Weight normalization normalizes the weight params so that the l2-norm of
  the matrix is equal to 1. This is implemented as a layer wrapper where
  each wrapped layer will have its params l2-normalized before computing
  its ``__call__`` output.
  
  Example usage::
  
    >>> import flax, flax.linen as nn
    >>> import jax, jax.numpy as jnp
  
    >>> class Baz(nn.Module):
    ...   @nn.compact
    ...   def __call__(self, x):
    ...     return nn.Dense(2)(x)
  
    >>> class Bar(nn.Module):
    ...   @nn.compact
    ...   def __call__(self, x):
    ...     x = Baz()(x)
    ...     x = nn.Dense(3)(x)
    ...     x = Baz()(x)
    ...     x = nn.Dense(3)(x)
    ...     return x
  
    >>> class Foo(nn.Module):
    ...   @nn.compact
    ...   def __call__(self, x):
    ...     x = nn.Dense(3)(x)
    ...     # l2-normalize all params of the second Dense layer
    ...     x = nn.WeightNorm(nn.Dense(4), variable_filter=None)(x)
    ...     x = nn.Dense(5)(x)
    ...     # l2-normalize all kernels in the Bar submodule and all params in the
    ...     # Baz submodule
    ...     x = nn.WeightNorm(Bar(), variable_filter={'kernel', 'Baz'})(x)
    ...     return x
  
    >>> # init
    >>> x = jnp.ones((1, 2))
    >>> model = Foo()
    >>> variables = model.init(jax.random.key(0), x)
    >>> flax.core.freeze(jax.tree_map(jnp.shape, variables))
    FrozenDict({
        params: {
            Bar_0: {
                Baz_0: {
                    Dense_0: {
                        bias: (2,),
                        kernel: (5, 2),
                    },
                },
                Baz_1: {
                    Dense_0: {
                        bias: (2,),
                        kernel: (3, 2),
                    },
                },
                Dense_0: {
                    bias: (3,),
                    kernel: (2, 3),
                },
                Dense_1: {
                    bias: (3,),
                    kernel: (2, 3),
                },
            },
            Dense_0: {
                bias: (3,),
                kernel: (2, 3),
            },
            Dense_1: {
                bias: (4,),
                kernel: (3, 4),
            },
            Dense_2: {
                bias: (5,),
                kernel: (4, 5),
            },
            WeightNorm_0: {
                Dense_1/bias/scale: (4,),
                Dense_1/kernel/scale: (4,),
            },
            WeightNorm_1: {
                Bar_0/Baz_0/Dense_0/bias/scale: (2,),
                Bar_0/Baz_0/Dense_0/kernel/scale: (2,),
                Bar_0/Baz_1/Dense_0/bias/scale: (2,),
                Bar_0/Baz_1/Dense_0/kernel/scale: (2,),
                Bar_0/Dense_0/kernel/scale: (3,),
                Bar_0/Dense_1/kernel/scale: (3,),
            },
        },
    })
  
  Attributes:
    layer_instance: Module instance that is wrapped with WeightNorm
    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).
    use_scale: If True, creates a learnable variable ``scale`` that is
      multiplied to the ``layer_instance`` variables after l2-normalization.
    scale_init: Initialization function for the scaling function.
    feature_axes: The feature axes dimension(s). The l2-norm is calculated by
      reducing the ``layer_instance`` variables over the remaining (non-feature)
      axes. Therefore a separate l2-norm value is calculated and a separate
      scale (if ``use_scale=True``) is learned for each specified feature. By
      default, the trailing dimension is treated as the feature axis.
    variable_filter: An optional iterable that contains string items. The
      WeightNorm layer will selectively apply l2-normalization to the
      ``layer_instance`` variables whose key path (delimited by '/') has a match
      with ``variable_filter``. For example, ``variable_filter={'kernel'}`` will
      only apply l2-normalization to variables whose key path contains 'kernel'.
      By default, ``variable_filter={'kernel'}``.
class_docstr_tok: ["L2","weight","normalization","(","https",":","\/\/arxiv.org\/pdf\/1602.07868.pdf",")",".","Weight","normalization","normalizes","the","weight","params","so","that","the","l2-norm","of","the","matrix","is","equal","to","1",".","This","is","implemented","as","a","layer","wrapper","where","each","wrapped","layer","will","have","its","params","l2-normalized","before","computing","its","``","__call__","``","output",".","Example","usage",":",":",">",">",">","import","flax",",","flax.linen","as","nn",">",">",">","import","jax",",","jax.numpy","as","jnp",">",">",">","class","Baz","(","nn.Module",")",":","...","@","nn.compact","...","def","__call__","(","self",",","x",")",":","...","return","nn.Dense","(","2",")","(","x",")",">",">",">","class","Bar","(","nn.Module",")",":","...","@","nn.compact","...","def","__call__","(","self",",","x",")",":","...","x","=","Baz","(",")","(","x",")","...","x","=","nn.Dense","(","3",")","(","x",")","...","x","=","Baz","(",")","(","x",")","...","x","=","nn.Dense","(","3",")","(","x",")","...","return","x",">",">",">","class","Foo","(","nn.Module",")",":","...","@","nn.compact","...","def","__call__","(","self",",","x",")",":","...","x","=","nn.Dense","(","3",")","(","x",")","...","#","l2-normalize","all","params","of","the","second","Dense","layer","...","x","=","nn.WeightNorm","(","nn.Dense","(","4",")",",","variable_filter=None",")","(","x",")","...","x","=","nn.Dense","(","5",")","(","x",")","...","#","l2-normalize","all","kernels","in","the","Bar","submodule","and","all","params","in","the","...","#","Baz","submodule","...","x","=","nn.WeightNorm","(","Bar","(",")",",","variable_filter=","{","'kernel","'",",","'Baz","'","}",")","(","x",")","...","return","x",">",">",">","#","init",">",">",">","x","=","jnp.ones","(","(","1",",","2",")",")",">",">",">","model","=","Foo","(",")",">",">",">","variables","=","model.init","(","jax.random.key","(","0",")",",","x",")",">",">",">","flax.core.freeze","(","jax.tree_map","(","jnp.shape",",","variables",")",")","FrozenDict","(","{","params",":","{","Bar_0",":","{","Baz_0",":","{","Dense_0",":","{","bias",":","(","2",",",")",",","kernel",":","(","5",",","2",")",",","}",",","}",",","Baz_1",":","{","Dense_0",":","{","bias",":","(","2",",",")",",","kernel",":","(","3",",","2",")",",","}",",","}",",","Dense_0",":","{","bias",":","(","3",",",")",",","kernel",":","(","2",",","3",")",",","}",",","Dense_1",":","{","bias",":","(","3",",",")",",","kernel",":","(","2",",","3",")",",","}",",","}",",","Dense_0",":","{","bias",":","(","3",",",")",",","kernel",":","(","2",",","3",")",",","}",",","Dense_1",":","{","bias",":","(","4",",",")",",","kernel",":","(","3",",","4",")",",","}",",","Dense_2",":","{","bias",":","(","5",",",")",",","kernel",":","(","4",",","5",")",",","}",",","WeightNorm_0",":","{","Dense_1\/bias\/scale",":","(","4",",",")",",","Dense_1\/kernel\/scale",":","(","4",",",")",",","}",",","WeightNorm_1",":","{","Bar_0\/Baz_0\/Dense_0\/bias\/scale",":","(","2",",",")",",","Bar_0\/Baz_0\/Dense_0\/kernel\/scale",":","(","2",",",")",",","Bar_0\/Baz_1\/Dense_0\/bias\/scale",":","(","2",",",")",",","Bar_0\/Baz_1\/Dense_0\/kernel\/scale",":","(","2",",",")",",","Bar_0\/Dense_0\/kernel\/scale",":","(","3",",",")",",","Bar_0\/Dense_1\/kernel\/scale",":","(","3",",",")",",","}",",","}",",","}",")","Attributes",":","layer_instance",":","Module","instance","that","is","wrapped","with","WeightNorm","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",")",".","use_scale",":","If","True",",","creates","a","learnable","variable","``","scale","``","that","is","multiplied","to","the","``","layer_instance","``","variables","after","l2-normalization",".","scale_init",":","Initialization","function","for","the","scaling","function",".","feature_axes",":","The","feature","axes","dimension","(","s",")",".","The","l2-norm","is","calculated","by","reducing","the","``","layer_instance","``","variables","over","the","remaining","(","non-feature",")","axes",".","Therefore","a","separate","l2-norm","value","is","calculated","and","a","separate","scale","(","if","``","use_scale=True","``",")","is","learned","for","each","specified","feature",".","By","default",",","the","trailing","dimension","is","treated","as","the","feature","axis",".","variable_filter",":","An","optional","iterable","that","contains","string","items",".","The","WeightNorm","layer","will","selectively","apply","l2-normalization","to","the","``","layer_instance","``","variables","whose","key","path","(","delimited","by","'\/","'",")","has","a","match","with","``","variable_filter","``",".","For","example",",","``","variable_filter=","{","'kernel","'","}","``","will","only","apply","l2-normalization","to","variables","whose","key","path","contains","'kernel","'",".","By","default",",","``","variable_filter=","{","'kernel","'","}","``","."]
docstr: Compute the l2-norm of the weights in ``self.layer_instance``
  and normalize the weights using this value before computing the
  ``__call__`` output.
  
  Args:
    *args: positional arguments to be passed into the call method of the
      underlying layer instance in ``self.layer_instance``.
    **kwargs: keyword arguments to be passed into the call method of the
      underlying layer instance in ``self.layer_instance``.
  
  Returns:
    Output of the layer using l2-normalized weights.
docstr_tok: ["Compute","the","l2-norm","of","the","weights","in","``","self.layer_instance","``","and","normalize","the","weights","using","this","value","before","computing","the","``","__call__","``","output",".","Args",":","*","args",":","positional","arguments","to","be","passed","into","the","call","method","of","the","underlying","layer","instance","in","``","self.layer_instance","``",".","*","*","kwargs",":","keyword","arguments","to","be","passed","into","the","call","method","of","the","underlying","layer","instance","in","``","self.layer_instance","``",".","Returns",":","Output","of","the","layer","using","l2-normalized","weights","."]
code: def __call__(self, *args, **kwargs):
      """Compute the l2-norm of the weights in ``self.layer_instance``
      and normalize the weights using this value before computing the
      ``__call__`` output.
  
      Args:
        *args: positional arguments to be passed into the call method of the
          underlying layer instance in ``self.layer_instance``.
        **kwargs: keyword arguments to be passed into the call method of the
          underlying layer instance in ``self.layer_instance``.
  
      Returns:
        Output of the layer using l2-normalized weights.
      """
  
      def layer_forward(layer_instance):
          return layer_instance(*args, **kwargs)
  
      return transforms.map_variables(
          layer_forward,
          trans_in_fn=lambda vs: jax.tree_util.tree_map_with_path(
              self._l2_normalize,
              vs,
          ),
          init=self.is_initializing(),
      )(self.layer_instance)
  
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 1 for calling the WeightNorm obj's underlying member method code with example usage: obj.__call__() and returns: transforms, layer_instance