Repository: google/jax

Issue:
core.call_bind aggressively raises args to top trace
```
def f(a_bool, y):
  if a_bool:
    return y + 1
  else:
    return y

jax.jit(jax.remat(f), static_argnums=0)(True, 1)
```

Results in:
```
ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected (in `bool`).
Use transformation parameters such as `static_argnums` for `jit` to avoid tracing input values.
See `https://jax.readthedocs.io/en/latest/faq.html#abstract-tracer-value-encountered-where-concrete-value-is-expected-error`.
Encountered value: Traced<ShapedArray(bool[], weak_type=True):JaxprTrace(level=-1/2)>
```

I think this arises from the full_raise occurring here when processing `remat_call_p` - which raise to the JaxprTrace when we told JIT we don't want to!
https://github.com/google/jax/blob/77901e9fa71f5b23066c70132a983ae57f655b39/jax/core.py#L1001

This also applies to user-defined call primitives using `core.call_bind`, resulting in unnecessary workarounds like this one in Haiku:
https://github.com/deepmind/dm-haiku/blob/49b21f7192dfdb3dc0a49cc097c8d3b0ccabb107/haiku/_src/named_call.py#L101-L109 


Discussion:
I think this might be working as intended. That is, `remat` itself will raise the abstraction level. You can use `jax.remat(f, concrete=True)` if you want to handle Python control flow inside a `remat`, but when mixing that together with `jit` it's possible to do redundant FLOPs in that case, so we wanted to make it opt-in. [This comment](https://github.com/google/jax/pull/1749#issuecomment-558267584) on #1749 has some detail.

WDYT?
Hmm I just tried with `concrete=True` and it didn't work... let's see...
I'm confused because calling `JaxprTrace.full_raise` should neither produce an 'unknown' tracer (see [1](https://github.com/google/jax/blob/77901e9fa71f5b23066c70132a983ae57f655b39/jax/interpreters/partial_eval.py#L98) [2](https://github.com/google/jax/blob/77901e9fa71f5b23066c70132a983ae57f655b39/jax/interpreters/partial_eval.py#L109)) nor raise the abstraction level. An unknown is only introduced when we call `JaxprTrace.instantiate_const`, and then it's only raised past the Concrete level if we call `JaxprTrace.instantiate_const_abstracted`.

Oh... I think I have an idea...
It's because [we stage out as much as possible](https://github.com/google/jax/blob/77901e9fa71f5b23066c70132a983ae57f655b39/jax/interpreters/partial_eval.py#L167) for calls. We really only want that behavior for `xla_call` and not other call primitives...

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -989,7 +989,7 @@ def process_env_traces(post_processor: str, primitive: Primitive,
   yield outs, tuple(todo)  # Ensure the aux output is immutable
 
 def _call_bind(processor: str, post_processor: str, primitive: Primitive,
-              f: lu.WrappedFun, *args, **params):
+               f: lu.WrappedFun, *args, **params):
   top_trace = find_top_trace(args)
   level = trace_state.trace_stack.next_level(True) if top_trace is None else top_trace.level
   params_tuple = tuple(params.items())
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -163,7 +163,8 @@ def default_process_primitive(self, primitive, tracers, params):
 
   def process_call(self, call_primitive, f: lu.WrappedFun, tracers, params):
     name = params.get('name', f.__name__)
-    if self.master.trace_type is StagingJaxprTrace:
+    if (self.master.trace_type is StagingJaxprTrace
+        and call_primitive in staged_out_calls):
       tracers = map(self.instantiate_const_abstracted, tracers)
     else:
       name = wrap_name(name, 'pe')
@@ -312,6 +313,7 @@ def _unmapped_aval(size, aval):
 
 custom_partial_eval_rules: Dict[core.Primitive, Callable] = {}
 call_partial_eval_rules: Dict[core.Primitive, Callable] = {}
+staged_out_calls: Set[core.Primitive] = set()
 
 
 def partial_eval(f, trace, pvs: Sequence[Optional[AbstractValue]], instantiate=False):
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -608,6 +608,7 @@ def _get_device(device, backend):
 xla_call = partial(core.call_bind, xla_call_p)
 xla_call_p.def_custom_bind(xla_call)
 xla_call_p.def_impl(_xla_call_impl)
+pe.staged_out_calls.add(xla_call_p)
 
 def _xla_call_translation_rule(c, axis_env,
                                in_nodes, name_stack, backend, name,


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1510,6 +1510,16 @@ def scanned_f(x, _):
 
     jax.grad(scan_bug)(1.0)  # doesn't crash
 
+  def test_remat_jit_static_argnum(self):
+    # https://github.com/google/jax/issues/2833
+    def f(a_bool, y):
+      if a_bool:
+        return y + 1
+      else:
+        return y
+
+    api.jit(api.remat(f, concrete=True), static_argnums=0)(True, 1)  # no crash
+
   def test_trivial_computations(self):
     x = np.array([1, 2, 3])
     y = api.jit(lambda x: x)(x)