torchscript_jit_trace() — pytorch Function Reference

Architecture documentation for the torchscript_jit_trace() function in common.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking calls 2 called by 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ torchscript_jit_trace() — pytorch Function Reference

Dependency Diagram

graph TD
  9b71719a_7134_fa38_fd95_b665398662db["torchscript_jit_trace()"]
  04a3a4a6_8db3_854d_a893_02c9542bf9dd["speedup_experiment()"]
  04a3a4a6_8db3_854d_a893_02c9542bf9dd -->|calls| 9b71719a_7134_fa38_fd95_b665398662db
  1eff8423_1f23_d138_6815_07d8dc29a749["_normalize_bench_inputs()"]
  9b71719a_7134_fa38_fd95_b665398662db -->|calls| 1eff8423_1f23_d138_6815_07d8dc29a749
  3df8de63_d71a_714c_93a1_1a0e24c0f362["load()"]
  9b71719a_7134_fa38_fd95_b665398662db -->|calls| 3df8de63_d71a_714c_93a1_1a0e24c0f362
  style 9b71719a_7134_fa38_fd95_b665398662db fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/common.py lines 1568–1575

def torchscript_jit_trace(model, example_inputs):
    optimized = JitTracedCache.load(model, example_inputs)

    def opt_jit_trace(_, example_inputs, collect_outputs=False):
        example_args, example_kwargs = _normalize_bench_inputs(example_inputs)
        return optimized(*example_args, **example_kwargs)

    return opt_jit_trace