docstr: Layer wrapper to decouple magnitude and direction of the layer's weights.
  
  This wrapper reparameterizes a layer by decoupling the weight's
  magnitude and direction. This speeds up convergence by improving the
  conditioning of the optimization problem. It has an optional data-dependent
  initialization scheme, in which initial values of weights are set as functions
  of the first minibatch of data. Both the weight normalization and data-
  dependent initialization are described in [Salimans and Kingma (2016)][1].
  
  Parameters
  ----------
  layer
      A `tf.keras.layers.Layer` instance. Supported layer types are
      `Dense`, `Conv2D`, and `Conv2DTranspose`. Layers with multiple inputs
      are not supported.
  data_init
      If `True` use data dependent variable initialization.
  **kwargs
      Additional keyword args passed to `tf.keras.layers.Wrapper`.
  
  Raises
  ------
  ValueError
      If `layer` is not a `tf.keras.layers.Layer` instance.
code: def __init__(self, layer, data_init: bool = True, **kwargs):
      """Layer wrapper to decouple magnitude and direction of the layer's weights.
  
      This wrapper reparameterizes a layer by decoupling the weight's
      magnitude and direction. This speeds up convergence by improving the
      conditioning of the optimization problem. It has an optional data-dependent
      initialization scheme, in which initial values of weights are set as functions
      of the first minibatch of data. Both the weight normalization and data-
      dependent initialization are described in [Salimans and Kingma (2016)][1].
  
      Parameters
      ----------
      layer
          A `tf.keras.layers.Layer` instance. Supported layer types are
          `Dense`, `Conv2D`, and `Conv2DTranspose`. Layers with multiple inputs
          are not supported.
      data_init
          If `True` use data dependent variable initialization.
      **kwargs
          Additional keyword args passed to `tf.keras.layers.Wrapper`.
  
      Raises
      ------
      ValueError
          If `layer` is not a `tf.keras.layers.Layer` instance.
      """
      if not isinstance(layer, tf.keras.layers.Layer):
          raise ValueError(
              "Please initialize `WeightNorm` layer with a `tf.keras.layers.Layer` "
              "instance. You passed: {input}".format(input=layer)
          )
  
      layer_type = type(layer).__name__
      if layer_type not in ["Dense", "Conv2D", "Conv2DTranspose"]:
          warnings.warn(
              "`WeightNorm` is tested only for `Dense`, `Conv2D`, and "
              "`Conv2DTranspose` layers. You passed a layer of type `{}`".format(
                  layer_type
              )
          )
  
      super(WeightNorm, self).__init__(layer, **kwargs)
  
      self.data_init = data_init
      self._track_trackable(layer, name="layer")
      self.filter_axis = -2 if layer_type == "Conv2DTranspose" else -1
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import numpy
  import warnings
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import logistic
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import quantized_distribution
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 to create a new WeightNorm object from inherited base classes: tf.keras.layers with example: obj = WeightNorm(layer, data_init)