torch.compiler.substitute_in_graph

torch.compiler.substitute_in_graph(original_fn, *, can_constant_fold_through=False, skip_signature_check=False) [исходный код]

Зарегистрировать обработчик-полифил для функции, обычно функции C из расширения C, который будет использоваться вместо исходной функции при встраивании исходной функции в граф.

Примечание

Обработчик-полифил используется только при встраивании исходной функции. Он не используется при прямом вызове исходной функции. В режиме eager декорированная функция вызывает производительную функцию C, а не обработчик-полифил.

Обработчик-полифил — это функция, которая вызывается вместо исходной функции при встраивании исходной функции. Обработчик-полифил должен иметь ту же сигнатуру и поведение, что и исходная функция.

Параметры:

• original_fn (callable) – Исходная функция, обычно функция C, для которой регистрируется обработчик-полифил.
• can_constant_fold_through (bool, optional) – Можно ли выполнять свёртку констант через обработчик-полифил. То есть если обработчик-полифил является чистой функцией и его аргументы — константы, результат обработчика-полифила можно свернуть в константу во время компиляции. По умолчанию False.
• skip_signature_check (bool, optional) – Следует ли пропустить проверку сигнатур исходной функции и обработчика-полифила. По умолчанию False.

Возвращает:

Декоратор, регистрирующий обработчик-полифил для исходной функции.

Тип возвращаемого значения:

Callable[[Callable[[~_P], _R]], Callable[[~_P], _R]]

Пример:

>>> import operator
>>> operator.indexOf([1, 2, 3, 4, 5], 3)
2
>>> torch.compile(operator.indexOf, fullgraph=True)([1, 2, 3, 4, 5], 3)
... # xdoctest: +SKIP("Long tracebacks")
Traceback (most recent call last):
...
torch._dynamo.exc.Unsupported: ...

>>> @torch.compiler.substitute_in_graph(operator.indexOf)
... def indexOf(a, b, /):
...     for i, item in enumerate(a):
...         if item is b or item == b:
...             return i
...     raise ValueError("sequence.index(x): x not in sequence")
>>>
>>> torch.compile(operator.indexOf, fullgraph=True)([1, 2, 3, 4, 5], 3)
2