Add static_argnames to pjit and lower

[rsepassi]

jit has static_argnames. Would be great to also have that for pjit and for the AoT api lower.

[andyehrenberg]

Would this be as simple as having

def extend_argnums_from_argnames(fun, argnums, argnames) -> Tuple[int, ...]:
    argnums, argnames = _infer_argnums_and_argnames(fun, argnums, argnames)
    if not argnames:
        return argnums
    else:
        try:
            signature = inspect.signature(fun)
        except ValueError:
            parameters: Mapping[str, inspect.Parameter] = {}
        else:
            parameters = signature.parameters
        added_argnums = tuple(
            i for i, (k, param) in enumerate(parameters.items())
            if param.kind == _POSITIONAL_OR_KEYWORD and k in argnames and i not in argnums
        )
        return argnums + added_argnums

And then within pjit, adding static_argnames as an argument and doing:

static_argnums = extend_argnums_from_argnames(fun, static_argnums, static_argnames)

Happy to write a PR if this seems like a reasonable approach!

[apaszke]

That looks approximately right, although I do think that it would be better to reuse _infer_argnums_and_argnames to be consistent with jit. If you have the time, PRs are definitely welcome!

[yashk2810]

This is fixed!