Additional input validation for transformations

[JeppeKlitgaard]

Adds additional input validation for *_argnames and *_argnums arguments to jax.jit.

Previously static_argnums and static_argnames could easily lead to silently dynamic arguments. In fact, two such cases were found in the jax source code using the new test coverage (c477ce3).

Additional test coverage has been added.

Includes a small refactoring in api.py which introduces _jit to reduce code duplication (this leads to a number of small changes in various tests). The refactoring has the additional benefit of making experimentation with #10476 easier.

For discussion see issue: #10601

Other places that could use additional validation (just add plumbing to validation logic introduced by this PR)

• jax.experiment.pjit
• jax.pmap
• jax.value_and_grad
• jax.custom_vjp
• jax.custom_jvp
• jax.hessian
• jax.jacrev
• jax.jacfwd
• jax.grad

I will add validation for the functions mentioned above in a separate PR once/if this one has been merged.

Handling of *args and **kwargs

In cases where variable positional and/or keyword arguments are used, we (rightly) assume that the function will know how to deal with these and thus it is safe to assume that for *args any argnum is valid. Similarly for **kwargs any argname is valid.

Importantly we still do some validation where possible. As an example, consider:

def f(a, /, b, *, c): ...

jit(f, static_argnames=("a",))   # This will fail since we know that `a` must be positional

jit(f, static_argnums=(2,))  # This will fail since `c` must be keyword

Additional improvements

argnums could be patched to contain positional-only arguments given as static_argnames using inspect.

Currently:

def f(a, /, b, *, c):
    print(a, b, c)

ff = jit(f, static_argnames=("a", "b", "c"))  # with validation disabled
ff(1, 2, c=3)
> Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)> 2 3
# Expected: 1 2 3

ff2 = jit(f, static_argnames=("b", "c"), static_argnums=(0,))
ff2(1, 2, c=3)
> 1 Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)> 3
# Expected 1 2 3

This is due to the way _infer_argnums_and_argnames works, which might be worthwhile to change in a separate PR. See further discussion in #10614

Fix: #10601
Fix: #10046

[hawkinsp]

This looks promising, thanks.

[JeppeKlitgaard]

@hawkinsp changes since last time are primarily support for negative argnums as per #10669

[YouJiacheng]

an idea:
If we do arg and kwargs inference using inspect tools following #10614 (comment) (there are some bugs in that comment), we might only need to catch error from bind? But this might need a breaking change...

[JeppeKlitgaard]

only

If I understand correctly, that would fall outside the scope of this PR. Further, it only addresses cases where bind is used (static) and not other argument annotations (donate).

This PR strictly deals with argument annotation validation at transformation time. I believe (but can't recall, and am traveling atm) that the binding process happens when the transformed function is called, not when the transformation is performed.

I am working on another PR that builds on top of this one which goes a little further, but this PR aims to be compatible and uncontroversial.

[YouJiacheng]

I don't mean that the bind for call, I mean that if we use Signature.bind_partial to inference argnums and argnames.
Okay I find it is hard to use Signature.bind_partial in current JAX API interface...

[hawkinsp]

This is looking good to me, but I have a final few nits.

[hawkinsp]

I think I'd leave this out of the list and add it locally in the test cases that explicitly want to test "without jit"? After all, it isn't a jit implementation.

[JeppeKlitgaard]

All the parameterised tests use the NOOP jitter, so for now I think it should remain. Adding it in locally for every parameterised test would just lead to a lot of code duplication without any benefit, I think. Can always refactor this later if there are any cases where the no-op test case is not used.

[JeppeKlitgaard]

Implemented the requested changes.

Outstanding nit: #10603 (review)

Let me know what you think @hawkinsp

[JeppeKlitgaard]

Squashed to one commit and fixed a test. Ready for review/merge