code: def test_batch_jacobian_larger_rank_and_dtype(self):
      w1 = tf.reshape(tf.range(24.0, dtype=tf.float64) + 1.0, (4, 2, 3))
      w2 = tf.reshape(
          tf.range(24.0, dtype=tf.float32) * 0.5 - 6.0, (4, 2, 1, 3)
      )
  
      def f(x, y):  # [4, 2, 3], [4, 2, 1, 3] -> [4, 3, 2]
          return tf.transpose(
              tf.cast(tf.math.cumsum(w1 * x, axis=-1), dtype=tf.float32)
              + tf.square(tf.reverse(w2 * y, axis=[-3]))[..., 0, :],
              perm=[0, 2, 1],
          )
  
      x = tf.cast(np.random.uniform(size=(4, 2, 3)), dtype=tf.float64)
      y = tf.cast(
          np.random.uniform(size=(4, 2, 1, 3)), dtype=tf.float32
      )
      jac = batch_jacobian(f, [x, y])
  
      # Check shapes.
      self.assertLen(jac, 2)
      self.assertAllEqual([4, 3, 2, 2, 3], jac[0].shape)
      self.assertAllEqual([4, 3, 2, 2, 1, 3], jac[1].shape)
      self.assertEqual(tf.float64, jac[0].dtype)
      self.assertEqual(tf.float32, jac[1].dtype)
  
      # Check results against `value_and_gradient`.
      out_shape = f(x, y).shape[1:]
      for i in range(np.prod(out_shape)):
          idx = (slice(None),) + np.unravel_index(i, out_shape)
          # pylint: disable=cell-var-from-loop
          _, grad = gradient.value_and_gradient(
              lambda x, y: f(x, y)[idx], [x, y]
          )
          print(grad[0].shape, jac[0].shape, jac[0][idx].shape)
          self.assertAllClose(grad[0], jac[0][idx])
          self.assertAllClose(grad[1], jac[1][idx])
  
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 13 for calling the GradientTest obj's underlying member method code with example usage: obj.test_batch_jacobian_larger_rank_and_dtype() and returns: tf