recovering static_argnums / static_argnames after jit

[femtomc]

Hi all!

I have transformations which I apply to functions -- that can occur after jax.jit. In other words, I jax.make_jaxpr functions, but I'm attempting to make it convenient to also specify static arguments through my transformations.

I haven't figured out a good way to do this, and I have many questions. For example:

@partial(jax.jit, static_argnames=["J"])
def some_fun(key, J):
    return some_comp(key, J)

If I later on attempt to jax.make_jaxpr, I want jax.make_jaxpr to be aware about the static indication.

Now, I may be doing something horrible (like, should I even be using jax.make_jaxpr on jitted functions?) -- but it seems like interpreters written on top of JAX likely need to solve this problem (or motivate a solution to it).

[jakevdp]

Good question! There isn't really any public API for this – the object returned by jax.jit is defined in C++ and does not directly expose the static_argnames or static_argnums parameter in its Python API – however, these components are part of the object's state used for pickling, so you can access them via f.__getstate__(). For example:

from jax import jit
from functools import partial

@partial(jit, static_argnames=['axis'])
def sum(x, axis=None):
  return x.sum(axis)

print(sum.__getstate__())

{'version': 1,
 'fun': <function __main__.sum(x, axis=None)>,
 'cache_miss': <function jax._src.api._cpp_jit.<locals>.cache_miss(*args, **kwargs)>,
 'get_device': <function jax._src.api._cpp_jit.<locals>.get_device_info()>,
 'has_explicit_device': False,
 'static_argnums': [1],
 'static_argnames': ['axis'],
 'donate_argnums': [],
 'cache': <jaxlib.xla_extension.jax_jit.CompiledFunctionCache at 0x7efbd64887f0>}

The returned value is a Python dict with entries related to what you're asking about (Note that these are versions of static_argnames/static_argnums that are processed based on the function's signature).

Be warned that if you use this, you should consider it a non-public API whose details may change without warning from release to release.