code: def test_batch_jacobian(self):
      w = tf.reshape(tf.range(12.0) + 1.0, (4, 3))
  
      def f(x):
          return tf.math.cumsum(w * x, axis=-1)
  
      self.assertAllEqual(
          tf.convert_to_tensor(
              [
                  [[1.0, 0.0, 0.0], [1.0, 2.0, 0.0], [1.0, 2.0, 3.0]],
                  [[4.0, 0.0, 0.0], [4.0, 5.0, 0.0], [4.0, 5.0, 6.0]],
                  [[7.0, 0.0, 0.0], [7.0, 8.0, 0.0], [7.0, 8.0, 9.0]],
                  [
                      [10.0, 0.0, 0.0],
                      [10.0, 11.0, 0.0],
                      [10.0, 11.0, 12.0],
                  ],
              ]
          ),
          batch_jacobian(f, tf.ones((4, 3))),
      )
  
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 12 for calling the GradientTest obj's underlying member method code with example usage: obj.test_batch_jacobian() and returns: tf