export() — pytorch Function Reference

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

Function python DynamoBenchmarks Microbenchmarking calls 2 called by 2

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ export() — pytorch Function Reference

Dependency Diagram

graph TD
  f258a26f_4fc8_d6ab_1167_c608fca925fb["export()"]
  3df8de63_d71a_714c_93a1_1a0e24c0f362["load()"]
  3df8de63_d71a_714c_93a1_1a0e24c0f362 -->|calls| f258a26f_4fc8_d6ab_1167_c608fca925fb
  d4bae6c8_efa9_4b5e_53ab_4e403ff2693e["load()"]
  d4bae6c8_efa9_4b5e_53ab_4e403ff2693e -->|calls| f258a26f_4fc8_d6ab_1167_c608fca925fb
  1eff8423_1f23_d138_6815_07d8dc29a749["_normalize_bench_inputs()"]
  f258a26f_4fc8_d6ab_1167_c608fca925fb -->|calls| 1eff8423_1f23_d138_6815_07d8dc29a749
  8964c48d_7b16_2455_878e_dbe471c4a037["_register_dataclass_output_as_pytree()"]
  f258a26f_4fc8_d6ab_1167_c608fca925fb -->|calls| 8964c48d_7b16_2455_878e_dbe471c4a037
  style f258a26f_4fc8_d6ab_1167_c608fca925fb fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/common.py lines 1491–1514

def export(model, example_inputs):
    from torch.export.dynamic_shapes import _combine_args, _tree_map_with_path

    example_args, example_kwargs = _normalize_bench_inputs(example_inputs)
    example_outputs = model(*example_args, **example_kwargs)
    _register_dataclass_output_as_pytree(example_outputs)

    combined_args = _combine_args(model, example_args, example_kwargs)
    dynamic_shapes = _tree_map_with_path(
        _produce_dynamic_shapes_for_export, combined_args
    )

    # NOTE: if args.export is ever enabled for --performance mode (rather than solely
    # --accuracy), we'll need to clone the model and subtract out extra memory usage, as
    # done in AOTInductorModelCache.
    ep = torch.export.export(
        model, example_args, example_kwargs, dynamic_shapes=dynamic_shapes, strict=True
    )

    def opt_export(_, example_inputs):
        example_args, example_kwargs = _normalize_bench_inputs(example_inputs)
        return ep.module()(*example_args, **example_kwargs)

    return opt_export