code: def testScanImpl(self, jit_scan, jit_f):
      rng = np.random.RandomState(0)
  
      d = rng.randn(2)
  
      def f(c, a):
          assert a.shape == (3,)
          assert c.shape == (4,)
          b = tf_np.cos(
              tf_np.sum(tf_np.sin(a))
              + tf_np.sum(tf_np.cos(c))
              + tf_np.sum(tf_np.tan(d))
          )
          c = tf_np.sin(c * b)
          assert (
              b.shape == ()
          )  # pylint: disable=g-explicit-bool-comparison
          return c, b
  
      if jit_f:
          f = extensions.jit(f)
  
      if jit_scan == "no_xla":
          scan = extensions.jit(extensions.scan, static_argnums=(0,))
      elif jit_scan == "xla_forced_compile":
          scan = extensions.jit(
              extensions.scan,
              static_argnums=(0,),
              xla_forced_compile=True,
          )
      else:
          scan = extensions.scan
  
      xs = rng.randn(5, 3)
      c = rng.randn(4)
  
      ans = scan(f, c, xs)
      expected = scan_reference(f, c, xs)
      if jit_scan == "xla_forced_compile":
          # xla.compile doesn't preserve list-vs-tuple properly for the outputs, so
          # we canonicalize them to lists here.
          expected = list(expected)
          ans = list(ans)
      self.assertDTypesEqual(expected, ans)
      self.assertAllClose(expected, ans)
  
imports: import functools
  import itertools
  from absl import flags
  from absl.testing import parameterized
  from jax import lax
  import numpy
  import tensorflow.compat.v2
  from trax.tf_numpy import extensions
  import trax.tf_numpy.numpy
label_desc: Use image node_id 22 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testScanImpl(jit_scan, jit_f) and returns: c,  b