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)
Tests that eval_on_shapes can be called within jit.
trax/tf_numpy/extensions/extensions_test.py