pytorch/benchmarks/dynamo/microbenchmarks/fx_microbenchmarks.py
Xuehai Pan c0ed38e644 [BE][Easy][3/19] enforce style for empty lines in import segments in benchmarks/ (#129754)
See https://github.com/pytorch/pytorch/pull/129751#issue-2380881501. Most changes are auto-generated by linter.

You can review these PRs via:

```bash
git diff --ignore-all-space --ignore-blank-lines HEAD~1
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/129754
Approved by: https://github.com/ezyang
2024-07-17 14:34:42 +00:00

32 lines
467 B
Python

import timeit
import torch.fx
N = 100000
K = 1000
def huge_graph():
def fn(x):
for _ in range(N):
x = x.sin()
return x
return torch.fx.symbolic_trace(fn)
def main():
g = huge_graph()
def fn():
for n in g.graph.nodes:
pass
t = min(timeit.repeat(fn, number=K, repeat=3))
print(f"iterating over {N*K} FX nodes took {t:.1f}s ({N*K/t:.0f} nodes/s)")
if __name__ == "__main__":
main()