make_jaxpr doesn't work with jitted function having static_argnums

[trancenoid]

I am learning JAX and was interested to see how JAX handles static arguments differently from normal arguments, but the following code doesn't seem to work. Am I missing something obvious or is this a bug?

from jax import make_jaxpr,jit
import jax
from functools import partial

x = 4
y = 2

@partial(jit, static_argnums = (1,))
def f(x, y):
  if y < 0:
    return x
  else:
    return x*2

# f(x,y) # works fine
make_jaxpr(f)(x,y) #throws error

Error :
ValueError: Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses. Static argument (index 1) of type <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'> for function f is non-hashable.

[jakevdp]

Hi, thanks for the question! Unlike jit, make_jaxpr does not have any mechanism to support static arguments directly, so the best way to proceed is to close over the static arguments, for example using partial:

from functools import partial
make_jaxpr(partial(f, y=y))(x)

[trancenoid]

Oh, gotcha!