code: def test_batching(self):
      x = tf.constant([[3.0, 4.0], [30.0, 40.0]])
      y = tf.constant([[5.0, 6.0], [50.0, 60.0]])
      z = tf.constant([[7.0, 8.0], [70.0, 80.0]])
      alpha = tf.constant(2.0)
      beta = tf.constant(1.0)
  
      with tf.GradientTape(persistent=True) as tape:
          tape.watch([alpha, beta])
  
          def body(i, state):
              x, y, z = state
              k = tf.cast(i + 1, tf.float32)
              return [
                  x * alpha - beta,
                  y * k * alpha * beta,
                  z * beta + x,
              ]
  
          out = for_loop(body, [x, y, z], [alpha, beta], 3)
      with self.subTest("independent_vars"):
          grad = tape.gradient(out[1], alpha)
          self.assertAllEqual(8712, grad)
      with self.subTest("dependent_vars"):
          grad = tape.gradient(out[2], beta)
          self.assertAllEqual(783, grad)
  
file_path: tf-quant-finance/tf_quant_finance/math/custom_loops_test.py
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
  from tf_quant_finance.math import custom_loops
label_desc: Use image node_id 4 for calling the ForLoopWithCustomGradientTest obj's underlying member method code with example usage: obj.test_batching() and returns: list