docstr: Tests that eval_on_shapes can be called within jit.
docstr_tok: ["Tests","that","eval_on_shapes","can","be","called","within","jit","."]
code: def testJitOfEvalOnShapes(self):
      """Tests that eval_on_shapes can be called within jit."""
  
      def transformer(f, **kwargs):
          def f_prime(*args):
              res = extensions.eval_on_shapes(f, **kwargs)(*args)
              return tf.nest.map_structure(
                  lambda x: tf_np.zeros(x.shape, x.dtype), res
              )
  
          return extensions.jit(
              f_prime, kwargs.get("static_argnums", ())
          )
  
      self._testEvalOnShapes(transformer, False)
  
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 14 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testJitOfEvalOnShapes() and returns: extensions, tf