def testSwitchBatched(self):

    def fun(index, x, y, z):
        branches = [(lambda xyz: xyz[0]), (lambda xyz: lax.neg(xyz[1])), (lambda xyz: lax.sign(xyz[2]))]
        return lax.switch(index, branches, (x, y, z))
    x = jnp.array(0)
    y = jnp.array([1, 2])
    z = jnp.array([3, 4])
    w = jnp.array(9)
    ans = api.vmap(fun, (None, 0, 0, None))(x, y, z, w)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, 0, None)))(x, y, z, w)
    expected = np.array([1, 2])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' not in str(jaxpr))
    x = jnp.array(1)
    ans = api.vmap(fun, (None, 0, 0, None))(x, y, z, w)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, 0, None)))(x, y, z, w)
    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, None))(x, y, z, w)
    expected = np.array([(- 3), (- 4)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    z = jnp.array(5)
    ans = api.vmap(fun, (None, 0, None, None))(x, y, z, w)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, None, None)))(x, y, z, w)
    expected = np.array([(- 5), (- 5)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' not in str(jaxpr))
    x = jnp.array([0, 1])
    ans = api.vmap(fun, (0, 0, None, None))(x, y, z, w)
    jaxpr = api.make_jaxpr(api.vmap(fun, (0, 0, None, None)))(x, y, z, w)
    expected = np.array([1, (- 5)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' in str(jaxpr))
    z = jnp.array([3, 4])
    w = jnp.array([9, 9])
    ans = api.vmap(fun)(x, y, z, w)
    jaxpr = api.make_jaxpr(api.vmap(fun))(x, y, z, w)
    expected = np.array([1, (- 4)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' in str(jaxpr))
def testSwitchBatched(self):

    def fun(index, x, y, z):
        branches = [(lambda xyz: xyz[0]), (lambda xyz: lax.neg(xyz[1])), (lambda xyz: lax.sign(xyz[2]))]
        return lax.switch(index, branches, (x, y, z))
    x = jnp.array(0)
    y = jnp.array([1, 2])
    z = jnp.array([3, 4])
    w = jnp.array(9)
    ans = api.vmap(fun, (None, 0, 0, None))(x, y, z, w)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, 0, None)))(x, y, z, w)
    expected = np.array([1, 2])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' not in str(jaxpr))
    x = jnp.array(1)
    ans = api.vmap(fun, (None, 0, 0, None))(x, y, z, w)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, 0, None)))(x, y, z, w)
    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, None))(x, y, z, w)
    expected = np.array([(- 3), (- 4)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    z = jnp.array(5)
    ans = api.vmap(fun, (None, 0, None, None))(x, y, z, w)
    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, None, None)))(x, y, z, w)
    expected = np.array([(- 5), (- 5)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' not in str(jaxpr))
    x = jnp.array([0, 1])
    ans = api.vmap(fun, (0, 0, None, None))(x, y, z, w)
    jaxpr = api.make_jaxpr(api.vmap(fun, (0, 0, None, None)))(x, y, z, w)
    expected = np.array([1, (- 5)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' in str(jaxpr))
    z = jnp.array([3, 4])
    w = jnp.array([9, 9])
    ans = api.vmap(fun)(x, y, z, w)
    jaxpr = api.make_jaxpr(api.vmap(fun))(x, y, z, w)
    expected = np.array([1, (- 4)])
    self.assertAllClose(ans, expected, check_dtypes=False)
    assert ('select' in str(jaxpr))