Our CPU CI doesn't have JAX and torchax so I'm installing it like this: https://github.com/pytorch/xla/pull/8781/files#diff-d30e144e1a94b9125c97915d4bd55eeb00d136257f96c9c753c255b66b1c00b4

What happens if we tie torch.aotautograd.grad() to here? won't work?

I guess we are skipping this op :)
This is a great exploration

• I'd love to understand how non-functional code behaves in this code path?
• How do we make tracing decisions on deta-dependent conditional operations?

@qihqi @tengyifei

What happens if we tie torch.aotautograd.grad() to here? won't work?

We'll have to option to either use this with AOTAutograd, or avoid AOTAutograd.

If we would like to use AOTAutograd, we can use AOTAutograd to turn a PyTorch function into FX Graph, then use torchax to turn the FX Graph to a JAX function, then use call_jax to turn the JAX function into a LazyTensor node.

If we would like to avoid AOTAutograd, we'll use torchax to turn a PyTorch function into a JAX function, use jax.grad to get the backward (also as a JAX function), then use call_jax to turn the JAX function into a LazyTensor node.

In the latter approach, we'll be able to use powerful JAX remat features like checkpoint_name (checkpoint or offload an arbitrary tensor).

I'd love to understand how non-functional code behaves in this code path?

In the general case, this won't support non-functional code. For example if some PyTorch function inserts a tensor into a global list, we'll just end up inserting a JAX tracer into the global list. But AOTAutograd has the same constraint. This limitation does not exist when only using LazyTensor because that framework blurs the boundary between tracing/execution at the cost of accidental graph breaks.

My hypothesis is that most high performance models (esp. ones in torchprime) won't do weird stuff like that.

We'll need to find some other solution to cover the long tail e.g. use dynamo.

How do we make tracing decisions on data-dependent conditional operations?

We can't do that (e.g. if (tensor) { ... }) in jax.jit. It'll raise a Python exception. To be clear, we also can't efficiently do that in LazyTensor as you'll graph break.

Similarly, my hypothesis is that high performance models won't do a data dependent branch. They need to rewrite the if into a jax.lax.cond.