Calling a jit'd function with or without kwarg name has different behavior

Consider this example:

def foo(x, foo_power=10):
   # ...

jit_foo = jax.jit(foo)
jit_foo(x, foo_power=11)
jit_foo(x, 11)

It would be reasonable to expect the two calls to jit_foo to have the same behavior/semantics. However, the first call with the kwarg name included will treat foo_power as a static arg, whereas the second call will treat it as a dynamic arg. This can cause errors if foo expects foo_power to be a compile-time constant (I ran into this jitting the fft method I’m working on).

A workaround is to always include kwargs in the static_argnums argument to jit. We might wanna consider doing this automatically in jit to have more predictable behavior by default.

[skye]

@mattjj, what’s unclear from the docstring (and confusing) is that the keyword vs positional args distinction is not made at the function definition, but rather at the call site. This makes sense from the implementation perspective, but IMO from the API perspective, it makes more sense to use how the args are specified in the definition since it can’t change across calls, and especially if you use jit() as a function decorator.

So I think the kwarg behavior is fine and makes sense, but we should use introspection to do it based on the function definition, not the function call.

[afrozenator]

+1 to @skye 's comment, also “principle of least astonishment”.