code: def test_aux_multi_arg(self):
      x = tf.constant([[2.0]])
      z = tf.constant([[3.0]])
  
      def f(x, z):
          return x**2 + z**2, (x, z)
  
      (y, aux), (dx, dz) = gradient.value_and_gradient(
          f, (x, z), has_aux=True
      )
  
      self.assertAllClose(x**2 + z**2, y)
      self.assertAllClose(2 * x, dx)
      self.assertAllClose(2 * z, dz)
      self.assertAllClose(x, aux[0])
      self.assertAllClose(z, aux[1])
  
      (dx, dz), aux = batch_jacobian(f, (x, z), has_aux=True)
  
      self.assertAllClose((2 * x)[..., tf.newaxis], dx)
      self.assertAllClose((2 * z)[..., tf.newaxis], dz)
      self.assertAllClose(x, aux[0])
      self.assertAllClose(z, aux[1])
  
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 15 for calling the GradientTest obj's underlying member method code with example usage: obj.test_aux_multi_arg() and returns: unknown,  x,  z