code: def test_simple_grad_wrt_initial_state(self):
      x = tf.constant([3.0])
      sigma = tf.constant(2.0)
  
      with tf.GradientTape() as tape:
          tape.watch(x)
  
          def body(i, state):
              del i
              x = state[0]
              return [x * sigma]
  
          out = for_loop(body, [x], [sigma], 3)[0]
  
      grad = tape.gradient(out, x)
      self.assertAllEqual([8], 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 2 for calling the ForLoopWithCustomGradientTest obj's underlying member method code with example usage: obj.test_simple_grad_wrt_initial_state() and returns: list