torch.fx.traceback.annotate_fn

torch.fx.traceback.annotate_fn(annotation_dict) [исходный код]

Декоратор, который оборачивает функцию в менеджер контекста annotate. Используйте его, если хотите аннотировать всю функцию, а не отдельный блок кода.

Примечание

Этот API не имеет обратной совместимости и может измениться в будущих выпусках.

Примечание

Этот API несовместим с fx.symbolic_trace или jit.trace. Он предназначен для использования с семейством трассировщиков PT2, например torch.export и dynamo.

Параметры:

annotation_dict (dict) – Словарь пользовательских пар ключ-значение, которые добавляются в метаданные трассировки FX для всех операций в функции.

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

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

Пример

Все операции в my_function будут содержать {“pp_stage”: 1} в своих метаданных.

>>> @annotate_fn({"pp_stage": 1})
... def my_function(x):
...     return x + 1

Предупреждение

Этот API является экспериментальным и НЕ имеет обратной совместимости.