Repository: google/jax
Issue: jnp.piecewise compiles on every call in some cases

Arose when submitting PR #7692. 

```
def g(x):
  one = jnp.ones_like(x)
  return jnp.where(x>0, one, -one)

def f1(x):
  one = jnp.ones_like(x)
  return jnp.piecewise(x, [x < 0, x > 0], [g, g, 0.])

def f2(x):
  return jnp.piecewise(x, [x < 0, x > 0], [-1., 1., 0.])

x = jnp.arange(-2, 3)
for i in range(5):
  print(i, 'f1', f1(x))

for i in range(5):
  print(i, 'f2', f2(x))
```

Checking the logs, I see:

```
DEBUG:absl:Compiling _cumulative_reduction (5315785472) for args (ShapedArray(bool[3,5]),).
DEBUG:absl:Compiling <unnamed wrapped function> (5316489152) for args (ShapedArray(bool[5]), ShapedArray(int32[]), ShapedArray(int32[])).
DEBUG:absl:Compiling <unnamed wrapped function> (5316499392) for args (ShapedArray(bool[5]), ShapedArray(int32[]), ShapedArray(int32[])).
DEBUG:absl:Compiling <lambda> (5316499904) for args (ShapedArray(int32[5]), ShapedArray(int32[])).
0 f1 [-1 -1  0  1  1]
DEBUG:absl:Compiling <unnamed wrapped function> (5316615552) for args (ShapedArray(bool[5]), ShapedArray(int32[]), ShapedArray(int32[])).
DEBUG:absl:Compiling <unnamed wrapped function> (5316614720) for args (ShapedArray(bool[5]), ShapedArray(int32[]), ShapedArray(int32[])).
1 f1 [-1 -1  0  1  1]
DEBUG:absl:Compiling <unnamed wrapped function> (5316651968) for args (ShapedArray(bool[5]), ShapedArray(int32[]), ShapedArray(int32[])).
DEBUG:absl:Compiling <unnamed wrapped function> (5316606656) for args (ShapedArray(bool[5]), ShapedArray(int32[]), ShapedArray(int32[])).
2 f1 [-1 -1  0  1  1]
DEBUG:absl:Compiling <unnamed wrapped function> (5316694400) for args (ShapedArray(bool[5]), ShapedArray(int32[]), ShapedArray(int32[])).
DEBUG:absl:Compiling <unnamed wrapped function> (5316706816) for args (ShapedArray(bool[5]), ShapedArray(int32[]), ShapedArray(int32[])).
3 f1 [-1 -1  0  1  1]
DEBUG:absl:Compiling <unnamed wrapped function> (5316745024) for args (ShapedArray(bool[5]), ShapedArray(int32[]), ShapedArray(int32[])).
DEBUG:absl:Compiling <unnamed wrapped function> (5316666944) for args (ShapedArray(bool[5]), ShapedArray(int32[]), ShapedArray(int32[])).
4 f1 [-1 -1  0  1  1]
0 f2 [-1 -1  0  1  1]
1 f2 [-1 -1  0  1  1]
2 f2 [-1 -1  0  1  1]
3 f2 [-1 -1  0  1  1]
4 f2 [-1 -1  0  1  1]
```
Every call to `f1` causes a compilation, whereas the equivalent `f2` does not.

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

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index ade57c14d..b6a9d3110 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -6478,6 +6478,16 @@ 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}")
+  consts = {i: c for i, c in enumerate(funclist) if not callable(c)}
+  funcs = {i: f for i, f in enumerate(funclist) if callable(f)}
+  return _piecewise(x, condlist, consts,
+                    frozenset(funcs.items()),  # dict is not hashable.
+                    *args, **kw)
+
+@partial(jit, static_argnames=['funcs'])
+def _piecewise(x, condlist, consts, funcs, *args, **kw):
+  funcs = dict(funcs)
+  funclist = [consts.get(i, funcs.get(i)) for i in range(len(condlist) + 1)]
   indices = argmax(cumsum(concatenate([zeros_like(condlist[:1]), condlist], 0), 0), 0)
   dtype = _dtype(x)
   def _call(f):
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 5275412e7..a8ea96212 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1826,6 +1826,15 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     # 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)
 
+  def testPiecewiseRecompile(self):
+    def g(x):
+      g.num_traces += 1
+      return x
+    g.num_traces = 0
+    x = jnp.arange(10.0)
+    for i in range(5):
+      jnp.piecewise(x, [x < 0], [g, 0.])
+    self.assertEqual(g.num_traces, 1)
 
   @parameterized.named_parameters(jtu.cases_from_list(
     {"testcase_name": "{}_perm={}_{}".format(