How to handle static args in custom Jaxpr interpreters?

[huterguier]

I've encountered another issue following this guide on writing custom Jaxpr interpreters. Whenever I use the transformed function within a jitted function that has static args, tracing fails since my transformation doesn't know which incoming arguments are marked as static by the outer jit.

import jax
from functools import wraps

def transform(fun):
    @wraps(fun)
    def wrapped(*args, **kwargs):
        closed_jaxpr = jax.make_jaxpr(fun)(*args, **kwargs)
        out = jax.core.eval_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.literals, args)
        return out
    return wrapped

def f(n):
    return jax.numpy.arange(n)

y = jax.jit(transform(f), static_argnums=0)(3)

jax.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[]

How can I determine which arguments are static? Is it enough to check which elements in args are hashable?
Thanks!

[jakevdp]

There's no absolutely general way to detect which arguments should be static and which should be dynamic: it's why jit requires the user to specify static arguments explicitly. My suggestion would be for you to do the same in your transform function; i.e.

def transform(fun, static_argnums=()):
  ...

and then split your arguments into static and dynamic, closing over static arguments and passing only the dynamic arguments to make_jaxpr. Does that make sense?

[huterguier]

Yes, I totally get what you mean! And I also thought about this option but it feels a bit unnecessary as I wouldn't do that for a regular function either. So there must be some way to determine/infer this during the trace right?
Can I assume all arguments that are not tracers to be static? I slightly modified the code above to illustrate what I mean:

import jax
from functools import wraps

def transform(fun):
    @wraps(fun)
    def wrapped(*args, **kwargs):
        print(args, kwargs)
        closed_jaxpr = jax.make_jaxpr(fun)(*args, **kwargs)
        out = jax.core.eval_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.literals, args)
        return out
    return wrapped

def f(x, n):
    return jax.numpy.arange(n) + x

y = jax.jit(transform(f), static_argnums=1)(1.0, 3)

(Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace>, 3) {}

In this case the first argument x is converted to a tracer as it is dynamic, while n is still a regular hashable integer.

[jakevdp]

I think that will work as long as your transformed function is only called under JIT, and none of the arguments are pytrees. If those assumptions are OK for your use-case, then this approach will be fine. If not, then you should have the user of transform specify static arguments explicitly.

[huterguier]

Thanks for your input! I went for the first approach and it seems to work as expected for now. I also made it work with pytrees by using a flattened version of the function as input to make_jaxpr. Using the transformation outside of JIT doesn't seem to be a problem either.
Here is a minimal working example in the same style as above:

import jax
import jax.core
from functools import wraps

def is_hashable(arg):
  try:
    hash(arg)
    return True
  except TypeError:
    return False

def flatten(fun, structure):
  def wrapped(*args_flat):
    args, kwargs = jax.tree.unflatten(structure, args_flat)
    out = fun(*args, **kwargs)
    return out
  return wrapped

def transform(fun):
    @wraps(fun)
    def wrapped(*args, **kwargs):
        args_flat, structure = jax.tree.flatten((args, kwargs))
        static_argnums = tuple(i for i, arg in enumerate(args_flat) if is_hashable(arg))
        closed_jaxpr, out_shape = jax.make_jaxpr(
          flatten(fun, structure),
          static_argnums=static_argnums, 
          return_shape=True
        )(*args_flat)
        out_structure = jax.tree.structure(out_shape)
        dynamic_args_flat = tuple(arg for arg in args_flat if not is_hashable(arg))
        out_flat = jax.core.eval_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.literals, *dynamic_args_flat)
        out = jax.tree_util.tree_unflatten(out_structure, out_flat)
        return out
    return wrapped

def f(t, n):
    return jax.numpy.arange(n) * t[0] + t[1]

y = jax.jit(transform(f), static_argnums=1)((2.0, 1.0), 3)

[jakevdp]

Yeah that looks OK – though be aware you may end up with different jaxprs inside and outside of JIT if you're passing non-array arguments (e.g. calling f(x=1), x will be hashable and static, but calling jax.jit(f)(x=1), x will be non-hashable and dynamic). That could cause important differences in the resulting execution path: for example, jnp.pow(x, y) lowers to lax.integer_pow when y is a static integer, and to lax.pow otherwise. But if you're OK with those kinds of differences between JIT-compiled and non-jit-compiled execution, this solution seems reasonable.