def test_disable_jit_cond_with_vmap(self):

    def fn(t):
        return lax.cond((t > 0), 0, (lambda x: 0), 0, (lambda x: 1))
    fn = api.vmap(fn)
    with api.disable_jit():
        _ = fn(jnp.array([1]))
def test_disable_jit_cond_with_vmap(self):

    def fn(t):
        return lax.cond((t > 0), 0, (lambda x: 0), 0, (lambda x: 1))
    fn = api.vmap(fn)
    with api.disable_jit():
        _ = fn(jnp.array([1]))