def METHOD_NAME(self):
    """Make sure the jit stub test case also works with jax."""
    try:
        import jax.numpy as np
        from jax import jit
        import jaxlib
    except ImportError:
        self.skipTest('jax is not available but required for this test.')
    class TestClass(object):
        @partial(jit, static_argnums=(0,))
        def f(self, x):
            return np.sqrt(x)
    x = np.linspace(0, 10, 11)
    result = TestClass().f(x)
    assert_near_equal(result**2, x)
    self.assertIsInstance(result, jaxlib.xla_extension.ArrayImpl)