Repository: google/jax

Issue:
Support static argnames/argnums to checkify
I tried using checkify today but it's not compatible with my code base because I'm relying on static argnames/argnums for some of the jitted functions. After wrapping the jitted functions with checkify, those inputs are misinterpreted as dynamic arguments and JAX raises an error because they aren't JAX types.


Discussion:
Thanks for raising this!

> those inputs are misinterpreted as dynamic arguments

Well, `static_argnums`/`static_argnames` only affects the particular `jit` call to which those arguments are passed. In particular, they don't change the meaning of the body. So if I understand correctly, this issue isn't checkify-specific: for any `f` where `jit(checkify(f), static_argnums=(0,))` would fail with the error you describe, I believe `jit(jit(f), static_argnums=(0,))` would fail in the same way. In both cases, the function `g` to which the `jit(g, static_argnums=(0,))` is applied only accepts pytrees of jax types (ie arrays) as arguments.

So, there are (at least) two possible problem statements here, and I just want to figure out which one you have in mind:
1. make both `jit(jit(f), static_argnums=(0,))` and `jit(checkify(f), static_argnums=(0,))` (and `scan` etc) Just Work, by changing the meaning of the outer `jit`'s `static_argnums` to affect the body as well; or
2. change `checkify` only so that this issue doesn't arise (e.g. maybe `checkify` should try to avoid assuming  _any_ inputs are jaxtypes, independent of any caller's `static_argnums`).

Did you have one of these in mind?
cc @LenaMartens 
I was thinking of your option 2, changing `checkify` to avoid assuming that all its inputs must be jaxtypes.

More generally, there could be a solution along the lines of this?

```python3
def with_static(transform, argname):
  specialized = {}
  def wrapped(fn, *args, **kwargs):
    value = kwargs.pop(argname, '_default')
    if value not in specialized:
      specialized[value] = transform(functools.partial(fn, **{argname: value}))
    return specialized[value](*args, **kwargs)
  return wrapped

fn = with_static(jit, 'string_arg')(fn)  # Same as jit(fn, static_argnames='string_arg')
fn = with_static(checkify, 'string_arg')(fn)
fn(array_arg=jnp.zeros(), string_arg='bar')
```
Actually, this line isnt quite right:
```python
fn = with_static(jit, 'string_arg')(fn)  # Same as jit(fn, static_argnames='string_arg')
```

`jit`'s `static_argnums`/`static_argnames` can't be implemented in user code (i.e. on top of a `jax.jit` API which doesn't itself have those options). If you try using the solution quoted here, you'll always get retraces/recompiles between two applications of `fn` even when passing the same value for `string_arg`, because you're creating a fresh callable object (the `partial` instance) every time. That's why `jit` has `static_argnums`/`static_argnames` built in: not as a convenience but because it's the only way to get the caching behavior we want.

This kind of caching approach can work better with `checkify`, though I'm a bit certain because it now operates differently from e.g. `jax.grad` (basically it has some internal caching which makes it a bit different). I'd have to think about it.

I guess a third option is:
3. add `static_argnums` / `static_argnames` to `checkify`, perhaps just as a convenience (though it's not as convenient as it Just Working).

I'd rather make it Just Work if we can!
I don't quite follow. The `functools.partial()` in my call only happens when the specialization value isn't already in the cache dictionary. So it would only retrace once per unique value for `string_arg`.
Sorry, I misread your code. (Actually I think it's a bit buggy in that you want to curry the `wrapped` fun on extra level and not get both `fn` and `args`/`kwargs` at the same time.)

You're right that you can get cache hits for equal values with this approach, so that part of my comment was mistaken. But now you have separate jit caches (so e.g. we can't get the same eviction logic we would get from having a single cache).

In general my only point was just that `with_static(jit, 'string_arg')(fn)` isn't _exactly_ the same as `jit(fn, static_argnames='string_arg')`.

Luckily, #15569 will solve the issue by making things Just Work! With that `checkify` will no longer place any constraints on the arguments passed to the `checkify`-decorated function. (Internally all the arguments are just closed over; we checked that the caching issues I was concerned about, basically the things #14250 and #14291 were fixing, no longer apply thanks to 6ec9082.)

Resolving patch:
diff --git a/jax/_src/checkify.py b/jax/_src/checkify.py
--- a/jax/_src/checkify.py
+++ b/jax/_src/checkify.py
@@ -382,20 +382,6 @@ def out_axes_thunk():
 def get_shaped_aval(val):
   return core.raise_to_shaped(core.get_aval(val))
 
-def initial_style_jaxpr(
-    fun: Callable, in_tree: PyTreeDef, in_avals: Sequence[core.AbstractValue]
-  ) -> Tuple[core.Jaxpr, List[Any], PyTreeDef]:
-  return _initial_style_jaxpr(fun, in_tree, tuple(in_avals))
-
-@weakref_lru_cache
-def _initial_style_jaxpr(fun, in_tree, in_avals):
-  # like control_flow._initial_style_jaxpr, but use flatten_fun not _nokwargs
-  fun_, out_tree = flatten_fun(lu.wrap_init(fun), in_tree)
-  debug = pe.debug_info(fun, in_tree, out_tree, False, 'checkify')
-  jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(fun_, in_avals, debug)
-  return jaxpr, consts, out_tree()
-
-
 def checkify_jaxpr(jaxpr: core.ClosedJaxpr, enabled_errors,
                    error: Error, *args) -> Tuple[Error, List[core.Value]]:
   err_vals, err_tree = jtu.tree_flatten(error)
@@ -1065,16 +1051,17 @@ def checkify(f: Callable[..., Out],
   """
   @traceback_util.api_boundary
   def checked_fun(*args, **kwargs):
+    # close over all arguments so they're not turned into abstract values.
+    in_tree = jtu.tree_structure(((), {}))
+    closed_f = lambda: f(*args, **kwargs)
     # stage:
-    flat_args, in_tree = tree_flatten((args, kwargs))
-    in_avals = map(get_shaped_aval, flat_args)
-    jaxpr_, consts, out_tree = initial_style_jaxpr(f, in_tree, in_avals)
+    fun_, out_tree = flatten_fun(lu.wrap_init(closed_f), in_tree)
+    debug = pe.debug_info(closed_f, in_tree, out_tree, False, 'checkify')
+    jaxpr_, _, consts = pe.trace_to_jaxpr_dynamic(fun_, (), debug)
     jaxpr = pe.close_jaxpr(pe.convert_constvars_jaxpr(jaxpr_))
     # checkify:
-    flat_args = jtu.tree_leaves((args, kwargs))
-    error, out_flat = checkify_jaxpr(jaxpr, errors, init_error,
-                                     *consts, *flat_args)
-    return error, jtu.tree_unflatten(out_tree, out_flat)
+    error, out_flat = checkify_jaxpr(jaxpr, errors, init_error, *consts)
+    return error, jtu.tree_unflatten(out_tree(), out_flat)
   return checked_fun
 
 def check(pred: Bool, msg: str, *fmt_args, **fmt_kwargs) -> None:


Test patch:
diff --git a/tests/checkify_test.py b/tests/checkify_test.py
--- a/tests/checkify_test.py
+++ b/tests/checkify_test.py
@@ -753,7 +753,7 @@ def f(x, y):
     cf = checkify.checkify(f, errors=checkify.nan_checks)
     err, _ = cf(jnp.array([-jnp.inf, 0, jnp.inf]), jnp.ones((3, 2)))
     self.assertIsNotNone(err.get())
-    self.assertStartsWith(err.get(), "nan generated by primitive: psum")
+    self.assertStartsWith(err.get(), "nan generated by primitive")
 
   def test_different_payload_effects(self):
     def f(x, y):
@@ -1214,6 +1214,18 @@ def test_fmt_args_array_type_error(self):
     trees_ok = lambda: checkify.check(False, "{}", {"hello": jnp.array(1.)})
     checkify.checkify(trees_ok)()
 
+  def test_checkify_non_jax_type_input(self):
+    _ = checkify.checkify(lambda x: 1.)("hi")  # does not crash
+
+  def test_checkify_static_args(self):
+    @checkify.checkify
+    def f(x):
+      if x:
+        return
+
+    _ = jax.jit(f, static_argnums=(0,))(True)
+
+
 class LowerableChecksTest(jtu.JaxTestCase):
   def setUp(self):
     super().setUp()