mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
022cd7c9aa
Add decorator `torch.compiler.substitute_in_graph` to register polyfill for unsupported C++ function to avoid graph break. This API provides an official way to add support for dynamo for third-party C extensions. Also, it can be used to simplify our implementation for `torch._dynamo.polyfill`.
5ee070266f/torch/_dynamo/variables/builtin.py (L97-L107)
Example:
```python
>>> import operator
>>> operator.indexOf([1, 2, 3, 4, 5], 3)
2
>>> torch.compile(operator.indexOf, fullgraph=True)([1, 2, 3, 4, 5], 3)
Unsupported: ...
>>> @torch.compiler.substitute_in_graph(operator.indexOf)
... def indexOf(sequence, x):
... for i, item in enumerate(sequence):
... if item is x or item == x:
... return i
... raise ValueError("sequence.index(x): x not in sequence")
>>> torch.compile(operator.indexOf, fullgraph=True)([1, 2, 3, 4, 5], 3)
2
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/133712
Approved by: https://github.com/jansel
26 lines
496 B
ReStructuredText
26 lines
496 B
ReStructuredText
.. currentmodule:: torch.compiler
|
|
|
|
.. automodule:: torch.compiler
|
|
|
|
.. _torch.compiler_api:
|
|
|
|
torch.compiler API reference
|
|
============================
|
|
|
|
For a quick overview of ``torch.compiler``, see :ref:`torch.compiler_overview`.
|
|
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
compile
|
|
reset
|
|
allow_in_graph
|
|
substitute_in_graph
|
|
assume_constant_result
|
|
list_backends
|
|
disable
|
|
cudagraph_mark_step_begin
|
|
is_compiling
|
|
is_dynamo_compiling
|