code: def test_shapes(self, state_dims, num_params, times):
      # Checks that the loop can handle various shapes and outputs correct shapes.
      def test_with_batch_shape(batch_shape):
          initial_state = [
              tf.ones(shape=batch_shape + (d,)) for d in state_dims
          ]
          params = [tf.constant(1.0) for _ in range(num_params)]
          with tf.GradientTape(persistent=True) as tape:
              tape.watch(initial_state)
              tape.watch(params)
  
              def body(i, state):
                  del i
                  if not params:
                      return state
                  sum_params = tf.add_n(params)
                  state = [s * sum_params for s in state]
                  return state
  
              final_state = for_loop(body, initial_state, params, times)
  
          for s_in in initial_state:
              for s_out in final_state:
                  grad = tape.gradient(s_out, s_in)
                  self.assertAllEqual(s_in.shape, grad.shape)
  
          for p in params:
              for s_out in final_state:
                  grad = tape.gradient(s_out, p)
                  self.assertAllEqual([], grad.shape)
  
      with self.subTest("no_batch"):
          test_with_batch_shape(batch_shape=())
      with self.subTest("simple_batch"):
          test_with_batch_shape(batch_shape=(5,))
      with self.subTest("complex_batch"):
          test_with_batch_shape(batch_shape=(2, 8, 3))
  
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 7 for calling the ForLoopWithCustomGradientTest obj's underlying member method code with example usage: obj.test_shapes(state_dims, num_params, times) and returns: state, state