Get composition structure from Jax composed function

[joeryjoery]

Perhaps this is more a Python question than a Jax question but is it possible with jax.jit, jax.vmap or other compositions to get something like a "closure resolution order" back from Python? And on top of that, how can I easily check whether a function f has been composed with Jax transformations like jit or grad, etc.?

For example, I would like to do something along the lines of this (but also possibly nested):

x = np.ones(2)

if isinstance(f, jax.vmap):
    y = f(x[None, ...])
else:
    y = f(x)

I tried doing this with inspect.getclosurevars(f), this gives me e.g., something like this:

import inspect

inspect.getclosurevars(jax.vmap(f))
>> ClosureVars(nonlocals={'fun': <function vmap.<locals>.vmap_f at 0x7fcd940e3af0>}, globals={'filtering_mode': <function filtering_mode at 0x7fcf2028a040>, 'is_under_reraiser': ...

It seems like this only reveals the outer-most vmap though if I wrap it multiple times. Also, it's still a bit awkward to check the fun element in ClosureVars whether vmap is in the string representation.

Is there a more general way to achieve this?

[jakevdp]

As far as I am aware, JAX has made no attempt to provide hooks for this level of inspection, rather JAX makes use of decorators/closures, and liberally uses Python's functools.wraps for this purpose. This means that wrapped functions do have a __wrapped__ attribute that point to the original function:

from jax import vmap, jit

def f(x):
  return x

f_mapped = vmap(f)
f_jitted = jit(f)
f_mapped_then_jitted = jit(f_mapped)

print(f_mapped.__wrapped__ is f)
# True
print(f_jitted.__wrapped__ is f)
# True
print(f_mapped_then_jitted.__wrapped__ is f_mapped)
# True
print(f_mapped_then_jitted.__wrapped__.__wrapped__ is f)
# True

That said, JAX does not add any attributes to tell you which transform has been used on the wrapped function. You may be able to use implementation details to infer this; for example:

>>> type(f)
function
>>> type(jax.vmap(f))
function
>>> type(jax.jit(f))                                                                                                                                            
jaxlib.xla_extension.CompiledFunction

But these are not considered public parts of the API, so I would not recommend relying on them.

I haven't thought about the pros and cons of implementing a user-facing API to track the transformation structure of functions, but you may wish to open a feature request if you have an application that would benefit from it.

[joeryjoery]

Thanks for the explanation, I was unaware of the __wrapped__ attribute. At least for what I want to do now, explicit checking with is should suffice. I think this could be helpful in general, for debugging and inspection for instance. In my, perhaps isolated, use-case I want to evaluate multiple varying functions on the same input, here checking if the function was vmapped would help me specify the correct array dimensions.

Another instance where this would be helpful is for preventing some exceptions or raising assertions, e.g., in jaxopt you are able to vmap over numerical solvers, but if you vmap then you must also jit the function being optimized -- otherwise, this would raise a ConcretizationError. Being able to check if vmapped could be helplful here.

So, it can have added value, the full extent thereof... I'm not entirely sure.

[YouJiacheng]

If you are not writing a library, I think you can wrap function transformations(vmap, jit etc.) provided by JAX, and use your wrapped transformations which will record the information you need instead.