Support static arguments for any function transformation

[danijar]

Following up on #15504, I ended up needing a more general solution because I kept running into the same issue with other function transformations that often don't support static arguments. My solution ended up being this this snippet that adds static argument support to any transformation. A meta transformation! :D

def static_support(transform):
  def new_transform(fun, *args, static=(), **kwargs):
    assert isinstance(static, (list, tuple)), static
    cache = {}
    def new_function(*args2, **kwargs2):
      sta = {k: v for k, v in kwargs2.items() if k in static}
      dyn = {k: v for k, v in kwargs2.items() if k not in static}
      key = hash(tuple(sta.get(n, '_default') for n in static))
      if key not in cache:
        specialized = bind(fun, **sta)
        specialized.__name__ = fun.__name__
        cache[key] = transform(specialized, *args, **kwargs)
      return cache[key](*args2, **dyn)
    return new_function
  return new_transform

fun = static_support(jit)(fun, static=['static_arg'])
fun = static_support(checkify)(fun, static=['static_arg'])
# ...

A simplifying design choice here is that static arguments are required to be passed as keyword arguments and cannot be passed as positional arguments, which I find much easier to specify. When the transformed function is already the output of some other transformation, counting argument names can be painful.

I'm wondering whether including a helper like the one above would be a useful addition, simplification, and unification for JAX? It might simplify the existing JAX transformations, provide a unified API to users, and allow users to use third-party transformations that may not support static arguments. Longer term, users could even be asked to apply the meta transformation themselves instead of repeating the API across the provided transformations.

[patrick-kidger]

You might find Equinox interesting, which does something similar: it has equinox.filter_{jit, grad, ...}, which automatically determines dynamic vs static arguments based on their type. After all, it mostly only makes sense to JIT wrt arrays, to grad wrt floating-point arrays, etc.

[danijar]

Thanks Patrick, I've heard of equinox and think it has a lot of interesting ideas. I use Ninjax.