class_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].
  
  #### Example
  
  ```python
    net = WeightNorm(tf_keras.layers.Conv2D(2, 2, activation='relu'),
           input_shape=(32, 32, 3), data_init=True)(x)
    net = WeightNorm(tf_keras.layers.Conv2DTranspose(16, 5, activation='relu'),
                     data_init=True)
    net = WeightNorm(tf_keras.layers.Dense(120, activation='relu'),
                     data_init=True)(net)
    net = WeightNorm(tf_keras.layers.Dense(num_classes),
                     data_init=True)(net)
  ```
  
  #### References
  
  [1]: Tim Salimans and Diederik P. Kingma. Weight Normalization: A Simple
       Reparameterization to Accelerate Training of Deep Neural Networks. In
       _30th Conference on Neural Information Processing Systems_, 2016.
       https://arxiv.org/abs/1602.07868
class_docstr_tok: ["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","]",".","#","#","#","#","Example","``","`","python","net","=","WeightNorm","(","tf_keras.layers.Conv2D","(","2",",","2",",","activation='relu","'",")",",","input_shape=","(","32",",","32",",","3",")",",","data_init=True",")","(","x",")","net","=","WeightNorm","(","tf_keras.layers.Conv2DTranspose","(","16",",","5",",","activation='relu","'",")",",","data_init=True",")","net","=","WeightNorm","(","tf_keras.layers.Dense","(","120",",","activation='relu","'",")",",","data_init=True",")","(","net",")","net","=","WeightNorm","(","tf_keras.layers.Dense","(","num_classes",")",",","data_init=True",")","(","net",")","``","`","#","#","#","#","References","[","1","]",":","Tim","Salimans","and","Diederik","P.","Kingma",".","Weight","Normalization",":","A","Simple","Reparameterization","to","Accelerate","Training","of","Deep","Neural","Networks",".","In","_30th","Conference","on","Neural","Information","Processing","Systems_",",","2016.","https",":","\/\/arxiv.org\/abs\/1602.07868"]
docstr: Initialize WeightNorm wrapper.
  
  Args:
    layer: A `tf_keras.layers.Layer` instance. Supported layer types are
      `Dense`, `Conv2D`, and `Conv2DTranspose`. Layers with multiple inputs
      are not supported.
    data_init: `bool`, 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.
docstr_tok: ["Initialize","WeightNorm","wrapper",".","Args",":","layer",":","A","`","tf_keras.layers.Layer","`","instance",".","Supported","layer","types","are","`","Dense","`",",","`","Conv2D","`",",","and","`","Conv2DTranspose","`",".","Layers","with","multiple","inputs","are","not","supported",".","data_init",":","`","bool","`",",","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=True, **kwargs):
      """Initialize WeightNorm wrapper.
  
      Args:
        layer: A `tf_keras.layers.Layer` instance. Supported layer types are
          `Dense`, `Conv2D`, and `Conv2DTranspose`. Layers with multiple inputs
          are not supported.
        data_init: `bool`, 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
  
file_path: probability/tensorflow_probability/python/layers/weight_norm.py
imports: import warnings
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import tf_keras
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)