Repository: google/jax
Issue: jnp.piecewise breaks on scalar inputs 

From documentation and analogy to onp.piecewise():
```python
import numpy as onp
onp.piecewise(onp.array(0.1), [onp.array(True), onp.array(False)], [-1, 1])
```
the jax's alternative should work on scalar inputs (null shape arrays)

```python
from jax import numpy as jnp
jnp.piecewise(jnp.array(0.1), [jnp.array(True), jnp.array(False)], [-1, 1])
```
However, this breaks with the following error on jax 0.2.8 and 0.2.10:

> TypeError                                 Traceback (most recent call last)
> <ipython-input-6-3886dc4a1876> in <module>
> ----> 1 jnp.piecewise(jnp.array(0.1), [jnp.array(True), jnp.array(False)], [-1, 1])
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in piecewise(x, condlist, funclist, *args, **kw)
>    5132   else:
>    5133     raise ValueError(f"with {nc} condition(s), either {nc} or {nc+1} functions are expected; got {nf}")
> -> 5134   indices = argmax(cumsum(vstack([zeros_like(condlist[:1]), condlist]), 0), 0)
>    5135   dtype = _dtype(x)
>    5136   def _call(f):
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in vstack(tup)
>    2741 @_wraps(np.vstack)
>    2742 def vstack(tup):
> -> 2743   return concatenate([atleast_2d(m) for m in tup], axis=0)
>    2744 row_stack = vstack
>    2745 
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in concatenate(arrays, axis)
>    2734   else:
>    2735     while len(arrays) > 1:
> -> 2736       arrays = [lax.concatenate(arrays[i:i+k], axis)
>    2737                 for i in range(0, len(arrays), k)]
>    2738     return arrays[0]
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in <listcomp>(.0)
>    2734   else:
>    2735     while len(arrays) > 1:
> -> 2736       arrays = [lax.concatenate(arrays[i:i+k], axis)
>    2737                 for i in range(0, len(arrays), k)]
>    2738     return arrays[0]
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/_src/lax/lax.py in concatenate(operands, dimension)
>     490     An array containing the concatenation.
>     491   """
> --> 492   return concatenate_p.bind(*operands, dimension=dimension)
>     493 
>     494 Precision = xla_client.PrecisionConfig.Precision
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/core.py in bind(self, *args, **params)
>     282     top_trace = find_top_trace(args)
>     283     tracers = map(top_trace.full_raise, args)
> --> 284     out = top_trace.process_primitive(self, tracers, params)
>     285     return map(full_lower, out) if self.multiple_results else full_lower(out)
>     286 
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/core.py in process_primitive(self, primitive, tracers, params)
>     620 
>     621   def process_primitive(self, primitive, tracers, params):
> --> 622     return primitive.impl(*tracers, **params)
>     623 
>     624   def process_call(self, primitive, f, tracers, params):
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
>     239 def apply_primitive(prim, *args, **params):
>     240   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
> --> 241   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)
>     242   return compiled_fun(*args)
>     243 
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/_src/util.py in wrapper(*args, **kwargs)
>     196         return f(*args, **kwargs)
>     197       else:
> --> 198         return cached(bool(config.x64_enabled), *args, **kwargs)
>     199 
>     200     wrapper.cache_clear = cached.cache_clear
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/_src/util.py in cached(_, *args, **kwargs)
>     189     @functools.lru_cache(max_size)
>     190     def cached(_, *args, **kwargs):
> --> 191       return f(*args, **kwargs)
>     192 
>     193     @functools.wraps(f)
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
>     264     return _xla_callable(lu.wrap_init(prim_fun), device, None, "prim", donated_invars,
>     265                          *arg_specs)
> --> 266   aval_out = prim.abstract_eval(*avals, **params)
>     267   if not prim.multiple_results:
>     268     handle_result = aval_to_result_handler(device, aval_out)
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/_src/lax/lax.py in standard_abstract_eval(prim, shape_rule, dtype_rule, weak_type_rule, *avals, **kwargs)
>    1992                          weak_type=weak_type)
>    1993   elif least_specialized is ShapedArray:
> -> 1994     return ShapedArray(shape_rule(*avals, **kwargs), dtype_rule(*avals, **kwargs),
>    1995                        weak_type=weak_type)
>    1996   elif least_specialized is UnshapedArray:
> 
> ~/opt/anaconda3/envs/grape_jax210/lib/python3.8/site-packages/jax/_src/lax/lax.py in _concatenate_shape_rule(*operands, **kwargs)
>    3425            "dimension {} for shapes {}.")
>    3426     shapes = [operand.shape for operand in operands]
> -> 3427     raise TypeError(msg.format(dimension, ", ".join(map(str, shapes))))
>    3428 
>    3429   concat_size = sum(o.shape[dimension] for o in operands)
> 
> TypeError: Cannot concatenate arrays with shapes that differ in dimensions other than the one being concatenated: concatenating along dimension 0 for shapes (1, 1), (1, 2).

Changed files:
jax/_src/numpy/lax_numpy.py, tests/lax_numpy_test.py, jax/test_util.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 20c7be0d8..1f0588fec 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -5141,12 +5141,12 @@ def piecewise(x, condlist, funclist, *args, **kw):
     funclist = [0] + list(funclist)
   else:
     raise ValueError(f"with {nc} condition(s), either {nc} or {nc+1} functions are expected; got {nf}")
-  indices = argmax(cumsum(vstack([zeros_like(condlist[:1]), condlist]), 0), 0)
+  indices = argmax(cumsum(concatenate([zeros_like(condlist[:1]), condlist], 0), 0), 0)
   dtype = _dtype(x)
   def _call(f):
     return lambda x: f(x, *args, **kw).astype(dtype)
   def _const(v):
-    return lambda x: full_like(x, v)
+    return lambda x: array(v, dtype=dtype)
   funclist = [_call(f) if callable(f) else _const(f) for f in funclist]
   return vectorize(lax.switch, excluded=(1,))(indices, funclist, x)
 
diff --git a/jax/test_util.py b/jax/test_util.py
index a5e9eb38b..7804a0f1c 100644
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -875,7 +875,7 @@ class JaxTestCase(parameterized.TestCase):
                               msg="Found\\n{}\\nExpecting\\n{}".format(what, expected))
 
   def _CompileAndCheck(self, fun, args_maker, *, check_dtypes=True,
-                       rtol=None, atol=None):
+                       rtol=None, atol=None, check_cache_misses=True):
     """Helper method for running JAX compilation and allclose assertions."""
     args = args_maker()
 
@@ -892,10 +892,11 @@ class JaxTestCase(parameterized.TestCase):
 
     cache_misses = xla.xla_primitive_callable.cache_info().misses
     python_ans = fun(*args)
-    self.assertEqual(
-        cache_misses, xla.xla_primitive_callable.cache_info().misses,
-        "Compilation detected during second call of {} in op-by-op "
-        "mode.".format(fun))
+    if check_cache_misses:
+      self.assertEqual(
+          cache_misses, xla.xla_primitive_callable.cache_info().misses,
+          "Compilation detected during second call of {} in op-by-op "
+          "mode.".format(fun))
 
     cfun = api.jit(wrapped_fun)
     python_should_be_executing = True
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 83458ef9a..1029a5dd0 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1647,17 +1647,18 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
        "shape": shape, "dtype": dtype, "ncond": ncond, "nfunc": nfunc}
       for ncond in [1, 2, 3]
       for nfunc in [ncond, ncond + 1]
-      for shape in nonempty_nonscalar_array_shapes
+      for shape in all_shapes
       for dtype in all_dtypes))
   def testPiecewise(self, shape, dtype, ncond, nfunc):
     rng = jtu.rand_default(self.rng())
     rng_bool = jtu.rand_int(self.rng(), 0, 2)
     funclist = [lambda x: x - 1, 1, lambda x: x, 0][:nfunc]
-    args_maker = lambda: (rng(shape, dtype), list(rng_bool((ncond,) + shape, bool)))
+    args_maker = lambda: (rng(shape, dtype), [rng_bool(shape, bool) for i in range(ncond)])
     np_fun = partial(np.piecewise, funclist=funclist)
     jnp_fun = partial(jnp.piecewise, funclist=funclist)
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=True)
-    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
+    # This is a higher-order function, so the cache miss check will fail.
+    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, check_cache_misses=False)
 
 
   @parameterized.named_parameters(jtu.cases_from_list(