From b6d40bac5a1c2b8980ab4533f53893a650a73e39 Mon Sep 17 00:00:00 2001 From: Boyuan Feng Date: Wed, 10 Apr 2024 21:31:22 -0700 Subject: [PATCH] Add runtime assert to non-strict export (#2226) Summary: This PR moves `insert_deferred_runtime_asserts` from dynamo to torch.fx.passes and uses it to add runtime assertion for non-strict export. X-link: https://github.com/pytorch/pytorch/pull/123681 Differential Revision: D55944267 --- .../dynamo/dynamobench/_dynamo/utils.py | 21 +------------------ 1 file changed, 1 insertion(+), 20 deletions(-) diff --git a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py b/userbenchmark/dynamo/dynamobench/_dynamo/utils.py index 809adc6b04..a7ce87277f 100644 --- a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py +++ b/userbenchmark/dynamo/dynamobench/_dynamo/utils.py @@ -97,6 +97,7 @@ from torch._dispatch.python import enable_python_dispatcher from torch._utils_internal import log_compilation_event +from torch.fx._utils import _format_graph_code, lazy_format_graph_code from torch.nn.modules.lazy import LazyModuleMixin from torch.utils._triton import has_triton, has_triton_package @@ -1996,26 +1997,6 @@ def tensor_always_has_static_shape( return False, None -def lazy_format_graph_code(name, gm, maybe_id=None): - def format_name(): - if maybe_id is not None: - return f"{name} {maybe_id}" - else: - return name - - return LazyString( - lambda: _format_graph_code( - f"===== {format_name()} =====\n", - gm.forward.__code__.co_filename, - gm.print_readable(print_output=False), - ) - ) - - -def _format_graph_code(name, filename, graph_str): - return f"TRACED GRAPH\n {name} {filename} {graph_str}\n" - - def lazy_format_graph_tabular(fn_name, gm): def inner(): try: