`jit` input validation can lead to silently dynamic variables

jax.jit (and presumable other similar transformations) currently does not validate static_argnames and static_argnums, which leads to silent failures in some cases. These can be particularly frustrating to debug and the additional input validation can be done with relatively little overhead.

A particularly headache inducing example of this could be when used in combination with positional-only (or keyword-only) arguments.

Examples:

def f(dyn, stat):
	print(dyn)
	print(stat)
    ...

def g(dyn, stat, /):
	print(dyn)
	print(stat)
    ...

jf = jit(f, static_argnames=("stat",)
jf(1, 2)
> Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
> 2

jg = jit(g, static_argnames=("stat",)
jg(1, 2)
> Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
> Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
# This is somewhat expected, but should at least raise an error

jit(f, static_argnames=("some_arg_that_does_not_exist",)
# No error!

jit(f, static_argnums=(9,)
# No error!

Solution

Input validation on jit and similar functions may be made stricter. This may require using inspect where possible.

I would be happy to try undertaking this.

[JeppeKlitgaard]

Would this not also apply to argnums?

I am working on argnum validation right now, and I think it already spotted an error in JAX source: https://github.com/google/jax/blob/ab7a60b3abdc1e4aec210d79926824b85082d2d7/jax/_src/lax/qdwh.py#L68-L69

I assume (and believe to recall) static_argnums is zero-indexed, so making argnum=3 will never make sense.

Could you confirm this, @hawkinsp?

Would the correct arg_nums be (0, 1, 2) or (1, 2) in this case?

[JeppeKlitgaard]

Those two cases nicely motivate stricter input validation, I think!

I will make a PR with the additional validation.