mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
86370fd658
Summary: A follow up of https://github.com/pytorch/pytorch/pull/150689. Test Plan: test_dynamo -k test_guard_filter_fn Differential Revision: D72722322 Pull Request resolved: https://github.com/pytorch/pytorch/pull/150936 Approved by: https://github.com/jansel
26 lines
868 B
Python
26 lines
868 B
Python
"""Hook system for Dynamo's guard functionality.
|
|
|
|
This module provides a way to register callback functions that are triggered during
|
|
guard-related operations.
|
|
|
|
The Hooks class manages two types of hook functions:
|
|
- guard_export_fn: Called when guards need to be exported, taking a GuardsSet as input
|
|
- guard_fail_fn: Called when a guard check fails, taking a GuardFail object as input
|
|
|
|
These hooks enable customization of guard export and failure handling behaviors.
|
|
"""
|
|
|
|
import dataclasses
|
|
from typing import Callable, Optional
|
|
|
|
from torch._guards import GuardsSet
|
|
|
|
from .types import GuardFail, GuardFilterEntry
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Hooks:
|
|
guard_export_fn: Optional[Callable[[GuardsSet], None]] = None
|
|
guard_fail_fn: Optional[Callable[[GuardFail], None]] = None
|
|
guard_filter_fn: Optional[Callable[[list[GuardFilterEntry]], list[bool]]] = None
|