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.


Discussion:
Marking `f1` with `@jit`makes the problem go away. So is this a bug? Unsure.
I think the problem is more why even jitted `g` got traced twice when `jnp.piecewise` is present.
Thanks for the report! The reason for the re-compilation when `piecewise` is not jitted is from these lines: https://github.com/google/jax/blob/d349086ca5c676f5bcbb26b07b2c8e5946598431/jax/_src/numpy/lax_numpy.py#L5807-L5808

Each function passed to `piecewise` is wrapped before passing it to `lax.switch` in order to pass args and kwargs.

I think the easiest fix for this would probably be to JIT-compile `jnp.piecewise` by default.
Hmm, turns out it's not so simple as that, because `funclist` can contain a mix of static and non-static arguments (either functions or possibly-traced constants). Still doable, but not as easy as adding `@jit`
I think the root cause of this is from the `switch` statement: each time `switch` encounters a function it re-traces it:
```python
In [1]: from jax.lax import switch

In [2]: def g(x):
   ...:   print("tracing", g, x)
   ...:   return x

In [3]: def f(x):
   ...:   return switch(0, [g], 1)

In [4]: f(1)
tracing <function g at 0x7feb1ee56ee0> 1
Out[4]: 1

In [5]: f(1)
tracing <function g at 0x7feb1ee56ee0> 1
Out[5]: 1
```
I'm not sure whether this is intended or not.
If you change `f` to
```python
def f(x):
  return switch(0, [g, g], 1)
```
subsequent calls don't retrace anything.
> Hmm, turns out it's not so simple as that, because `funclist` can contain a mix of static and non-static arguments (either functions or possibly-traced constants). Still doable, but not as easy as adding `@jit`

I may be oversimplifying the situation, but finding a workaround seems rather easy. Changing
```python
  funclist = [_call(f) if callable(f) else _const(f) for f in funclist]
```
to
```python
  d = {}  
  funclist = [d.setdefault(f, _call(f)) if callable(f) else _const(f) for f in funclist]
```
works, with the added benefit/(unexpected behaviour?) that even not jitted function don't get trace twice. Consider

```python
@jax.jit
def g(x):
  print('g')
  return x

def f(x):
  return jnp.piecewise(x, [x < 0], [g, g])

x = jnp.r_[1, 1, -1, -1, 0]
f(x)
```
cc @froystig for what seems like potentially missed caching in `lax.switch`
It's not tracing the single branch at all, it's [simply calling it](https://github.com/google/jax/blob/bbfd8f7cfc540d7a93b9c8e7b4460fef4af9f560/jax/_src/lax/control_flow.py#L636-L637) in line. That's a correct optimization, and I think a desirable one.

In what sense is this a root cause?
Revisiting this. With better understanding of JAX now, I think the root cause of the issue is that `jnp.piecewise` is creating temporary functions... https://github.com/google/jax/blob/10af170a85a20b86b44d6d026e9ef1669e9af0ce/jax/_src/numpy/lax_numpy.py#L6489-L6494
...and passing them to `lax.switch`, which traces the functions it receives... https://github.com/google/jax/blob/7fa6b1b5fafaf22503bc194ba66a00a6859bdadd/jax/_src/lax/control_flow.py#L656-L657
... and since all trace cacheing is keyed on the function ID, temporary functions are essentially never cached. So the `jnp.piecewise` implementation is similar to the second block here:
```python
import jax.numpy as jnp
from jax import lax

def g(x):
  print(f'tracing g({x})')
  return x

def anon(func):
    return lambda *args: func(*args)

# g is traced once here:
lax.switch(0, [g, g], 1)
# tracing g(Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>)

# g is traced multiple times here:
lax.switch(0, [anon(g), anon(g)], 1)
# tracing g(Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>)
# tracing g(Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>)
```
I think the best solution here would be to make `jnp.piecewise` a wrapper for an underlying JIT-compiled function that separates out the static and non-static arguments it receives, so that the entire `lax.switch` call would be JIT compiled and appropriately cached.

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- 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):


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1826,6 +1826,15 @@ def testPiecewise(self, shape, dtype, ncond, nfunc):
     # 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(