Repository navigation
ENH: testing.lazy_xp_function: torch.compile support? #664
Description
Hey @crusaderky @rgommers , how feasible is it to add support for torch.compile to xpx.lazy_xp_function, such that we could have a torch_compile=True like jax_jit=True in SciPy?
cc @steppi perhaps you have also thought about this.
Activity
lucascolley commentedon Mar 31, 2026
steppi commentedon Mar 31, 2026
It's something I've been planning to look into but I haven't gotten to it yet. Thanks for taking the initiative on this.
crusaderky commentedon Mar 31, 2026
I am utterly ignorant of how torch.compile works, so I cannot answer this question without substantial analysis.
lucascolley commentedon Apr 2, 2026
Had a quick chat with claude and it was optimistic that pretty much the same approach as
jax_autojitwould work, usingtorch.utils._pytree.register_pytree_nodeinstead.lucascolley commentedon Apr 3, 2026
I opened a draft PR at gh-668, would be grateful if anyone could take a look.