def testCondGradVmapNan(self):
    eps = 0.001

    def safe1(x):
        return lax.cond((x < eps), (lambda _: eps), (lambda _: jnp.sqrt(x)), ())
    out = api.grad((lambda x: api.vmap(safe1)(x).sum()))(np.zeros(10))
    self.assertFalse(np.isnan(out).any())
def testCondGradVmapNan(self):
    eps = 0.001

    def safe1(x):
        return lax.cond((x < eps), (lambda _: eps), (lambda _: jnp.sqrt(x)), ())
    out = api.grad((lambda x: api.vmap(safe1)(x).sum()))(np.zeros(10))
    self.assertFalse(np.isnan(out).any())