[Question] jax.vmap & nested jax.vmap behavior

[ZaberKo]

Hello,

Thanks for your great work. I have a question about the behavior of torch2jax wrapped function on jax.vmap. I see it uses ffi.ffi_call("torch_call", outshapes, vmap_method="sequential"). Does it mean that the vmap call will be called sequentially in pytorch? If so, is there any way to improve it by using other vmap_method options?

[rdyro]

Thanks! Yes, your understanding is correct, it'll rewrite the vmap in terms of a sequential loop.

Unfortunately, there isn't a way around this at the moment as far as I can tell because PyTorch assumes sequential calls to itself and I get numerical errors in tests when calling PyTorch in parallel.

However, if in your specific application you get correct numerics with parallel calls to PyTorch, we could enable this as a configurable option in torch2jax. What do you think? I can either implement this as a configurable option or I'd welcome a PR if you're up for it (if in your case PyTorch handles parallel calls correctly).

[ZaberKo]

Thanks! Yes, your understanding is correct, it'll rewrite the vmap in terms of a sequential loop.

Unfortunately, there isn't a way around this at the moment as far as I can tell because PyTorch assumes sequential calls to itself and I get numerical errors in tests when calling PyTorch in parallel.

However, if in your specific application you get correct numerics with parallel calls to PyTorch, we could enable this as a configurable option in torch2jax. What do you think? I can either implement this as a configurable option or I'd welcome a PR if you're up for it (if in your case PyTorch handles parallel calls correctly).

Unfortunately, I'm not very familiar with the low-level coding in pytorch or jax. Maybe I can't be of much help with a PR 😂.

In my user case, the pytorch function is natively vectorizable. According to jax.pure_callback, vmap_method="expand_dims" or vmap_method="broadcast_all" will automatically handle nested jax vmaps on a flatten new dimension. Not sure if this is also applied to ffi_call. If true, there should be only one call to the pytorch function, where no parallel calls to pytorch will happen.

Besides the option vmap_method, I think has_side_effect can also be added as a configurable option.

[rdyro]

Let me run some numerical experiments and I can get back to you!

[rdyro]

@ZaberKo

I added experimental support for specifying vmap_method to both torch2jax and torch2jax_with_vjp, if you could provide feedback on if it helps, that'd be amazing!

torch2jax/torch2jax/api.py

Line 45 in b11ddb1

vmap_method: str = "sequential",

On a small test case, it is indeed batching the pytorch function!

(I'll leave has_side_effect for another time)

I'll keep this issue open until we can confirm that there are no corner cases where the non-default ("sequential") vmap_method breaks something.

[ZaberKo]

Great! I will try it and give some feedbacks