code: def test_accumulating_for_loop_grap_param(self, params_test):
      # Multiple  number of iterations produce correct gradients for params
      x = np.arange(24)
      x = np.reshape(x, [4, 3, 2])
      sigma_np = 2.0
      initial_state = tf.convert_to_tensor(x, dtype=tf.float64)
      sigma = tf.convert_to_tensor(sigma_np, dtype=tf.float64)
  
      def fn(initial_state, sigma):
          def body(i, state):
              del i
              x = state[0]
              return [x * sigma]
  
          return for_loop(body, [initial_state], [sigma], [3, 5])[0]
  
      # Tests for parameters gradient
      if params_test:
  
          def fwd_grad_fn(sigma):
              g = lambda sigma: fn(initial_state, sigma)
              return tff.math.fwd_gradient(
                  g, sigma, use_gradient_tape=True
              )
  
          def value_and_gradient(sigma):
              g = lambda sigma: fn(initial_state, sigma)
              return tff.math.value_and_gradient(
                  g, sigma, use_gradient_tape=True
              )
  
          expected_val = np.stack(
              [sigma_np**3 * x, sigma_np**5 * x], axis=0
          )
          expected_fwd_grad = np.stack(
              [3 * sigma_np**2 * x, 5 * sigma_np**4 * x], axis=0
          )
          expected_grad = np.sum(expected_fwd_grad)
          with self.subTest("ParamsForwardGradXLA"):
              fwd_grad = tf.function(fwd_grad_fn, jit_compile=True)(
                  sigma
              )
              self.assertAllClose(fwd_grad, expected_fwd_grad)
          with self.subTest("ParamsValueAndGradXLA"):
              val, grad = tf.function(
                  value_and_gradient, jit_compile=True
              )(sigma)
              self.assertAllClose(expected_val, val)
              self.assertAllClose(grad, expected_grad)
          with self.subTest("ParamsForwardGrad"):
              fwd_grad = fwd_grad_fn(sigma)
              self.assertAllClose(fwd_grad, expected_fwd_grad)
          with self.subTest("ParamsValueAndGrad"):
              val, grad = value_and_gradient(sigma)
              self.assertAllClose(expected_val, val)
              self.assertAllClose(grad, expected_grad)
      # Tests for initial state gradient
      if not params_test:
  
          def state_fwd_grad_fn(initial_state):
              g = lambda initial_state: fn(initial_state, sigma)
              return tff.math.fwd_gradient(
                  g, initial_state, use_gradient_tape=True
              )
  
          def state_value_and_gradient(initial_state):
              g = lambda initial_state: fn(initial_state, sigma)
              return tff.math.value_and_gradient(
                  g, initial_state, use_gradient_tape=True
              )
  
          expected_val = np.stack(
              [sigma_np**3 * x, sigma_np**5 * x], axis=0
          )
          expected_fwd_grad = np.stack(
              [
                  sigma_np**3 * np.ones_like(x),
                  sigma_np**5 * np.ones_like(x),
              ],
              axis=0,
          )
          expected_grad = np.sum(expected_fwd_grad, axis=0)
          with self.subTest("StateForwardGradXLA"):
              fwd_grad = tf.function(
                  state_fwd_grad_fn, jit_compile=True
              )(initial_state)
              self.assertAllClose(fwd_grad, expected_fwd_grad)
          with self.subTest("StateValueAndGradXLA"):
              val, grad = tf.function(
                  state_value_and_gradient, jit_compile=True
              )(initial_state)
              self.assertAllClose(expected_val, val)
              self.assertAllClose(grad, expected_grad)
          with self.subTest("StateForwardGrad"):
              fwd_grad = state_fwd_grad_fn(initial_state)
              self.assertAllClose(fwd_grad, expected_fwd_grad)
          with self.subTest("StateValueAndGrad"):
              val, grad = state_value_and_gradient(initial_state)
              self.assertAllClose(expected_val, val)
              self.assertAllClose(grad, expected_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 8 for calling the ForLoopWithCustomGradientTest obj's underlying member method code with example usage: obj.test_accumulating_for_loop_grap_param(params_test) and returns: for_loop, list, tff, tff, tff, tff