code: def test_get_output_spec_loop(self):
      args = (
          np.array(0.0, dtype=np.float64),
          (
              tf.convert_to_tensor(0.0),
              tf.convert_to_tensor([1.0, 1.0], dtype=tf.float64),
          ),
      )
      additional_args = (
          tf.convert_to_tensor([[3], [4]], dtype=tf.int32),
      )
      # Trace using both positional and keyword args.
      results = callable_util.get_output_spec(
          _return_args_from_infinite_loop,
          *args,
          additional_loop_vars=additional_args,
      )
      self.assertAllEqualNested(
          tf.nest.map_structure(
              lambda x: tf.convert_to_tensor(x).shape,
              args + additional_args,
          ),
          tf.nest.map_structure(lambda x: x.shape, results),
      )
      self.assertAllAssertsNested(
          self.assertEqual,
          tf.nest.map_structure(
              lambda x: tf.convert_to_tensor(x).dtype,
              args + additional_args,
          ),
          tf.nest.map_structure(lambda x: x.dtype, results),
      )
  
file_path: probability/tensorflow_probability/python/internal/callable_util_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import callable_util
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the CallableUtilTest obj's underlying member method code with example usage: obj.test_get_output_spec_loop() without return types