vmap for a class instance or member funtions

[Qiustander]

Hi all, I am curious about the implementation of vmap function. I want to use it to do auto-batching.

My case is like:

def batch_fn(my_class):
    def inner(parameter):
        my_class.update_paras(parameter)
        return another_fn(my_class)
    return inner

jit(vmap(batch_fn(my_class)))

As can be seen, my_class.update_paras is a member function that would modify some attributes in the class. Is vmap a multiple-threads function? Would it cause synchronization problem? What is the safe usage of vmap? Thanks!

[jakevdp]

Hi - thanks for the question! JAX transformations like vmap, jit, grad, etc. require pure function semantics (see JAX Sharp Bits: Pure Functions), and so functions like the one you mention are not going to be compatible.

It's hard to say more from the stripped-down example that you give, but I'd try re-expressing your function in a way that avoids implicit mutation: for example you could register my_class as a pytree (see Extending Pytrees) and have your function return the updated version of my_class.

[patrick-kidger]

If you want to do "classes with JAX", I'd recommend giving Equinox a try.

+1 to Jake's comment about avoiding mutation, as this is often a footgun when used with JAX transformations. (And something like Equinox will do its best to guide you along the happy path.)