pure_callback for custom class instances

[Qottmann]

Hello,

Is there a way to tell jax.jit to "ignore" the instantiation of a class object in a jit-compilation pipeline?

I am looking for something similar to using pure_callback but for a CustomClass object. I.e. I have the following situation. Creating the CustomClass instance outside the to-be-jitted function and then calling its method inside works:

# This works
class_instance = CustomClass(*arg, **kwargs)

@jax.jit
def func(x):
    return class_instance.do_computation(x)

I want to instantiate class_instance inside func. However, creating the class instance inside the function call does not work anymore because the class instantiation performs some illegal python control flow incompatible with jit-compiling. Changing CustomClass is not really an option, it is coming from an external library that I dont want to temper with.

# This does not work
@jax.jit
def func(x):
    class_instance = CustomClass(*arg, **kwargs)
    return class_instance.do_computation(x)

So my idea is to wrap the instantiation by class_instance = jax.pure_callback(CustomClass, result_shape_dtypes, *args, **kwargs) - however, I dont know what to put for result_shape_dtypes since CustomClass does not have a dtype (And more generally I guess pure_callback is not intended to do something like this.

Is this the right approach for this situation? Any pointers to possible solutions or workaround would be appreciated!

[jakevdp]

pure_callback, by design, can return jax-comparible objects, and CustomClass is not a jax-compatible object. Given the mechanics of pure_callback, I doubt this requirement will be relaxed.

My primary suggestion would be to modify CustomClass so that it is compatible with traced inputs, then you should have no problem using it in a jit-compiled function.

If this is not possible, then perhaps you could instead do a pure callback around a function that both creates and calls the class? i.e. something like this:

def g(x):
  class_instance = CustomClass(*arg, **kwargs)
  return class_instance.do_computation(x)

@jax.jit
def func(x):
  ...
  return jax.pure_callback(g, result_shape_dtype, x)

Without more details it's hard to guess what the best approach would be. What do you think?

[Qottmann]

Thank you for the quick reply @jakevdp !

I think making CustomClass compatible with traced inputs is not an immediate possibility because it is from a third party.
The suggested workaround works in principle, but the problem is that now .do_computation(x) is not jit-compiled anymore, but that is what I want to achieve. The odd thing to me is that when class_instance is created outside the jit scope it works, so in principle jax-jit knows how to use black-box functions. I would want to declare the instantiation a black-box operation for jax-jit to ignore 🤔 .

A bit hacky, but is there a way to make the instantiation of CustomClass jax-compatible to use pure_callback (or something similar)? E.g. creating a new class that inherits from CustomClass and has a shape and dtype attribute defined?

[jakevdp]

It's hard to guess what might work without more information. Can you edit your question to add a minimal reproducible example?

[patrick-kidger]

You might be able to try equinox.filter_pure_callback, which is equivalent to jax.pure_callback but handles arbitrary Python objects.

[vaishnkv]

Can you give a toy example (on the usage of the function )? I couldn't get the usage of the argument named result_shape_dtypes for the equinox.filter_pure_callback function. I have gone through the documentation.

Suppose I have a function , that takes input a jnp.array and returns str type obj, Can I use filter_pure_callback for this ?

The below description is given in the documentation
result_shape_dtypes: A PyTree specifying the output of callback. It should have a jax.ShapeDtypeStruct in place of any JAX arrays.

[Qottmann]

Sorry for not answering, we in the end went with another library and did not have that problem anymore. But found @patrick-kidger 's answer to help and work on another instance and marked as answer.