kwargs sometimes cannot work well with jit

Minimal repro:

from functools import partial
from jax import jit

@partial(jit, static_argnums=(3,))
def f(a, b, c, d):
  if d > 0:
    return a + b - c
  else:
    return d

f(a=1, b=2, c=3, d=4) would complain about TypeError: Jitted function has static_argnums=(3,) but was called with only 0 positional arguments.

but f(1, 2, 3, 4) works.

[mattjj]

This error is intentional, and unfortunately (AFAICT) necessary.

The underlying issue is that in Python there’s no way to reliably determine a function’s signature (it can be changed or hidden by decorators, functools.partial, etc.), which means there’s no way to match named arguments to their positions (and vice-versa) given a calling set of args/kwargs. That in turn means static_argnums doesn’t work with names. See #595 for a bit more info.

One solution is to use a wrapper like this:

from functools import partial
from jax import jit

@partial(jit, static_argnums=(3,))
def _f(a, b, c, d):
  if d > 0:
    return a + b - c
  else:
    return d

def f(a, b, c, d):
  return _f(a, b, c, d)

Now callers of f can write f(1, 2, 3, 4) or f(a=1, b=2, c=3, d=4) or whatever calling convention they like; the single layer of indirection from f to _f means that Python standardizes the arguments for us. We use this trick inside JAX in a lot of places, like in random.py but also for every lax primitive, which is one reason why you see things like def sin(x): return sin_p.bind(x) in lax.py.

WDYT?

[jakevdp]

The original issue has now been fixed for jit, and similar improvements for vmap and other transforms are tracked in #10614