get_shapes() — pytorch Function Reference

Architecture documentation for the get_shapes() function in kernels.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking called by 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ get_shapes() — pytorch Function Reference

Dependency Diagram

graph TD
  5fc0e2cf_eb29_ca36_919f_3d0adc7f386e["get_shapes()"]
  3f9404a0_b5ca_d039_2c4e_873e66d4608d["benchmark()"]
  3f9404a0_b5ca_d039_2c4e_873e66d4608d -->|calls| 5fc0e2cf_eb29_ca36_919f_3d0adc7f386e
  style 5fc0e2cf_eb29_ca36_919f_3d0adc7f386e fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/genai_layers/kernels.py lines 458–468

    def get_shapes(self) -> tuple[tuple[int, ...], ...]:
        # TODO: OOM for (32768, 65536) on h100
        return (
            (32768, 256),
            (32768, 512),
            (32768, 1024),
            (32768, 2048),
            (32768, 4096),
            (32768, 8192),
            (32768, 16384),
        ) + extra_shapes_for_norm