ENH: testing.lazy_xp_function: torch.compile support

[lucascolley]

Closes gh-664

@rgommers this was easier than I expected. Is there someone familiar with Dynamo that we could ping? A few things that came up:

• The following deprecation warning with a confusing message given that we are already using torch.compile:

E           torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
E           DeprecationWarning: `torch.jit.script_method` is not supported in Python 3.14+ and may break. Please switch to `torch.compile` or `torch.export`.

• Some failing tests here with:

E                       torch._dynamo.exc.FailOnRecompileLimitHit: recompile_limit reached with fullgraph=True. Excessive recompilations can degrade performance due to the compilation overhead of each recompilation. To monitor recompilations, enable TORCH_LOGS=recompiles. If recompilations are expected, consider increasing torch._dynamo.config.cache_size_limit to an appropriate value.

(Draft PR as this will need docs updates before merge.)

[lucascolley]

claude reckoned that, unlike JAX, we needn't treat iterables with this special case for torch.compile

[lucascolley]

(minor: could reduce some LoC perhaps)

[lucascolley]

@ev-br FYI

[lucascolley]

I think the warnings about to flood CI should be fixed by data-apis/array-api-compat#411 (comment).

[lucascolley]

ready for review

[j-bowhay]

@lucascolley what's the current status here? Is it just review bandwidth preventing this from moving forward?

[lucascolley]

@lucascolley what's the current status here? Is it just review bandwidth preventing this from moving forward?

Yes, pretty much. I had also noted in the linked issue:

Currently unclear how to deal with accepting a variable amount of graph breaks / recompilations.

[ev-br]

Currently unclear how to deal with accepting a variable amount of graph breaks / recompilations.

While I haven't written a single LOC, a reasonable design looks similar what is done for dask: an extra num_graphbreaks : int parameter (a made-up name), where a zero valueal maps onto fullgraph=True.

Also we definitely want to test things twice: in the eager mode and under torch dynamo. Not sure if it's already done for Jax (and if yes, where) or if it needs implementing.

[j-bowhay]

@lucascolley what's the current status here? Is it just review bandwidth preventing this from moving forward?

Yes, pretty much. I had also noted in the linked issue:

Currently unclear how to deal with accepting a variable amount of graph breaks / recompilations.

Do we need that level of granularity? Could we not have 3 options:

• No support for torch.compile
• Support for torch.compile but with full_graph=False
• Support for torch.compile but with full_graph=True

[lucascolley]

@lucascolley what's the current status here? Is it just review bandwidth preventing this from moving forward?

Yes, pretty much. I had also noted in the linked issue:

Currently unclear how to deal with accepting a variable amount of graph breaks / recompilations.

Do we need that level of granularity? Could we not have 3 options:

* No support for `torch.compile`

* Support for `torch.compile` but with `full_graph=False`

* Support for `torch.compile` but with `full_graph=True`

Would be happy to start with that at least if we can get that working.

[j-bowhay]

Do you have a preference for what that api would look like? Would it be acceptable to have both a torch_compile and torch_full_graph argument?

[lucascolley]

Do you have a preference for what that api would look like? Would it be acceptable to have both a torch_compile and torch_full_graph argument?

not sure yet what I think is nicest... see also #917

[j-bowhay]

Do you have a preference for what that api would look like? Would it be acceptable to have both a torch_compile and torch_full_graph argument?

not sure yet what I think is nicest... see also #917

I guess the one difference being we also want to be able to turn compile on/off and control if graph breaks are allowed

[lucascolley]

CI is happy now, with fullgraph=False. TODOs:

• ENH: fix is_jax_array under torch.compile array-api-compat#464
• investigate the warnings that are ignored here
• figure out the API for fullgraph=True
• add all three modes (eager, no fullgraph, fullgraph) to array-api-extra's CI matrix
• investigate recompilations / graph breaks in array-api-extra and perhaps remove some of them
• add docs
• try to plumb this into SciPy