xref: /aosp_15_r20/external/pytorch/torch/utils/benchmark/examples/simple_timeit.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# mypy: allow-untyped-defs
2"""Trivial use of Timer API:
3
4$ python -m examples.simple_timeit
5"""
6
7import torch
8
9import torch.utils.benchmark as benchmark_utils
10
11
12def main():
13    timer = benchmark_utils.Timer(
14        stmt="x + y",
15        globals={"x": torch.ones((4, 8)), "y": torch.ones((1, 8))},
16        label="Broadcasting add (4x8)",
17    )
18
19    for i in range(3):
20        print(f"Run: {i}\n{'-' * 40}")
21        print(f"timeit:\n{timer.timeit(10000)}\n")
22        print(f"autorange:\n{timer.blocked_autorange()}\n\n")
23
24
25if __name__ == "__main__":
26    main()
27