docstr: Tests that gradients flow through the auxiliary output.
docstr_tok: ["Tests","that","gradients","flow","through","the","auxiliary","output","."]
code: def test_aux_grads(self):
      """Tests that gradients flow through the auxiliary output."""
      x = tf.constant([[2.0]])
  
      def f(x):
          return x**2, x**2
  
      def f2(x):
          (_, aux), _ = gradient.value_and_gradient(f, x, has_aux=True)
          return aux
  
      y, dx = gradient.value_and_gradient(f2, x)
      self.assertAllClose(x**2, y)
      self.assertAllClose(2 * x, dx)
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.gradient import batch_jacobian
  from tensorflow_probability.python.math.gradient import value_and_gradient_with_auto_expansion
label_desc: Use image node_id 16 for calling the GradientTest obj's underlying member method code with example usage: obj.test_aux_grads() and returns: unknown,  unknown, aux