def testCondBatched(self):

    def fun(x, y, z):
        pred = lax.lt(x, 3)
        true_fun = (lambda y: y)
        false_fun = (lambda z: lax.neg(z))
        return lax.cond(pred, y, true_fun, z, false_fun)
    x = jnp.array(2)
    y = jnp.array([1, 2])
    z = jnp.array([3, 4])
    ans = api.vmap(fun, (None, 0, 0))(x, y, z)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, 0)))(x, y, z)
    expected = np.array([1, 2])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' not in str(jaxpr))
    x = jnp.array(4)
    ans = api.vmap(fun, (None, 0, 0))(x, y, z)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, 0)))(x, y, z)
    expected = np.array([(- 3), (- 4)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' not in str(jaxpr))
    fun = api.jit(fun)
    ans = api.vmap(fun, (None, 0, 0))(x, y, z)
    expected = np.array([(- 3), (- 4)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    z = jnp.array(5)
    ans = api.vmap(fun, (None, 0, None))(x, y, z)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, None)))(x, y, z)
    expected = np.array([(- 5), (- 5)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' not in str(jaxpr))
    x = jnp.array([2, 4])
    ans = api.vmap(fun, (0, 0, None))(x, y, z)
    jaxpr = api.make_jaxpr(api.vmap(fun, (0, 0, None)))(x, y, z)
    expected = np.array([1, (- 5)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' in str(jaxpr))
    z = jnp.array([3, 4])
    ans = api.vmap(fun)(x, y, z)
    jaxpr = api.make_jaxpr(api.vmap(fun))(x, y, z)
    expected = np.array([1, (- 4)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' in str(jaxpr))
def testCondBatched(self):

    def fun(x, y, z):
        pred = lax.lt(x, 3)
        true_fun = (lambda y: y)
        false_fun = (lambda z: lax.neg(z))
        return lax.cond(pred, y, true_fun, z, false_fun)
    x = jnp.array(2)
    y = jnp.array([1, 2])
    z = jnp.array([3, 4])
    ans = api.vmap(fun, (None, 0, 0))(x, y, z)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, 0)))(x, y, z)
    expected = np.array([1, 2])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' not in str(jaxpr))
    x = jnp.array(4)
    ans = api.vmap(fun, (None, 0, 0))(x, y, z)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, 0)))(x, y, z)
    expected = np.array([(- 3), (- 4)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' not in str(jaxpr))
    fun = api.jit(fun)
    ans = api.vmap(fun, (None, 0, 0))(x, y, z)
    expected = np.array([(- 3), (- 4)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    z = jnp.array(5)
    ans = api.vmap(fun, (None, 0, None))(x, y, z)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, None)))(x, y, z)
    expected = np.array([(- 5), (- 5)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' not in str(jaxpr))
    x = jnp.array([2, 4])
    ans = api.vmap(fun, (0, 0, None))(x, y, z)
    jaxpr = api.make_jaxpr(api.vmap(fun, (0, 0, None)))(x, y, z)
    expected = np.array([1, (- 5)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' in str(jaxpr))
    z = jnp.array([3, 4])
    ans = api.vmap(fun)(x, y, z)
    jaxpr = api.make_jaxpr(api.vmap(fun))(x, y, z)
    expected = np.array([1, (- 4)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' in str(jaxpr))