hooks.py 891 B

1234567891011121314151617181920212223242526
  1. """Hook system for Dynamo's guard functionality.
  2. This module provides a way to register callback functions that are triggered during
  3. guard-related operations.
  4. The Hooks class manages two types of hook functions:
  5. - guard_export_fn: Called when guards need to be exported, taking a GuardsSet as input
  6. - guard_fail_fn: Called when a guard check fails, taking a GuardFail object as input
  7. These hooks enable customization of guard export and failure handling behaviors.
  8. """
  9. import dataclasses
  10. from collections.abc import Callable, Sequence
  11. from torch._guards import GuardsSet
  12. from .types import GuardFail, GuardFilterEntry
  13. @dataclasses.dataclass
  14. class Hooks:
  15. guard_export_fn: Callable[[GuardsSet], None] | None = None
  16. guard_fail_fn: Callable[[GuardFail], None] | None = None
  17. guard_filter_fn: Callable[[Sequence[GuardFilterEntry]], Sequence[bool]] | None = (
  18. None
  19. )