code: def testGradientValues(self, model_type):
      model = self._define_model(
          model_type, self.data_dim, self.num_hidden
      )
      model.build(self.random_input.shape)
      self.evaluate([v.initializer for v in model.weights])
      norm_layer = self._get_maybe_norm_layer(model)
      self.evaluate(model(self.random_input))
  
      # Calculate gradients of a loss function with respect to `g` and `v`
      with tf.GradientTape(persistent=True) as tape:
          out = model(self.random_input)
          loss = tf.norm(out)
      grad_loss_v = self.evaluate(tape.gradient(loss, norm_layer.v))
      grad_loss_g = self.evaluate(tape.gradient(loss, norm_layer.g))
  
      # Build an identical model with no weight normalization
      model_no_norm = self._define_model(
          model_type,
          self.data_dim,
          self.num_hidden,
          use_weight_norm=False,
      )
      model_no_norm.build(self.random_input.shape)
      self.evaluate([v.initializer for v in model_no_norm.weights])
      layer = self._get_maybe_norm_layer(model_no_norm)
      self.evaluate(model_no_norm(self.random_input))
  
      # Calculate gradients of the loss function with respect to the kernel
      kernel = (
          tf.nn.l2_normalize(
              norm_layer.v, axis=norm_layer.kernel_norm_axes
          )
          * norm_layer.g
      )
      with tf.GradientTape() as tape:
          with tf.control_dependencies(
              [
                  layer.kernel.assign(kernel),
                  layer.bias.assign(norm_layer.layer.bias),
              ]
          ):
              out_no_norm = model_no_norm(self.random_input)
              loss_no_norm = tf.norm(out_no_norm)
      grad_loss_w = self.evaluate(
          tape.gradient(loss_no_norm, layer.kernel)
      )
  
      # Check that gradient equations (3) in Salimans and Kingma (2016) hold.
      g = self.evaluate(norm_layer.g)
      v = self.evaluate(norm_layer.v)
      v_norm = self.evaluate(
          tf.sqrt(
              tf.reduce_sum(
                  tf.square(v), axis=norm_layer.kernel_norm_axes
              )
          )
      )
  
      grad_loss_v_calculated = (
          g / v_norm
      ) * grad_loss_w - g * grad_loss_g * v / tf.square(v_norm)
      grad_loss_g_calculated = np.diag(
          np.matmul(np.transpose(grad_loss_w), v) / v_norm
      )
  
      self.assertAllClose(grad_loss_g, grad_loss_g_calculated)
      self.assertAllClose(grad_loss_v, grad_loss_v_calculated)
  
imports: import tempfile
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.layers import weight_norm
label_desc: Use image node_id 10 for calling the WeightNormTest obj's underlying member method code with example usage: obj.testGradientValues(model_type) without return types