xref: /aosp_15_r20/external/pytorch/torch/utils/benchmark/utils/_stubs.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1from typing import Any, Callable, Dict, Protocol, runtime_checkable
2
3
4class TimerClass(Protocol):
5    """This is the portion of the `timeit.Timer` API used by benchmark utils."""
6    def __init__(
7        self,
8        stmt: str,
9        setup: str,
10        timer: Callable[[], float],
11        globals: Dict[str, Any],
12        **kwargs: Any,
13    ) -> None:
14        ...
15
16    def timeit(self, number: int) -> float:
17        ...
18
19
20@runtime_checkable
21class TimeitModuleType(Protocol):
22    """Modules generated from `timeit_template.cpp`."""
23    def timeit(self, number: int) -> float:
24        ...
25
26
27class CallgrindModuleType(Protocol):
28    """Replicates the valgrind endpoints in `torch._C`.
29
30    These bindings are used to collect Callgrind profiles on earlier versions
31    of PyTorch and will eventually be removed.
32    """
33    __file__: str
34    __name__: str
35
36    def _valgrind_supported_platform(self) -> bool:
37        ...
38
39    def _valgrind_toggle(self) -> None:
40        ...
41