def _get_interpreter_name_for_var(var):

frame = inspect.currentframe()

if not frame:

raise RuntimeError("failed to inspect frame")

i = 0

while i < frames_up + 1:

frame = frame.f_back

if not frame:

raise RuntimeError("failed to get frame")

i += 1

f_locals = frame.f_locals

for k, v in f_locals.items():

if isinstance(v, torch.Tensor) and var is v:

return k if k != "self" else ""

return ""

# since Parameter.detach() always creates a new torch.Tensor instance,

# id(v) doesn't work with it. So we always get the Parameter or Buffer

# as values, and deduplicate the params using Parameters and Buffers

state_dict = module.state_dict(keep_vars=True)

filtered_dict = type(state_dict)()

seen_ids: set[int] = set()

for k, v in state_dict.items():

if id(v) in seen_ids:

continue

seen_ids.add(id(v))

if keep_vars:

filtered_dict[k] = v

else:

filtered_dict[k] = v.detach()

def __init__(

self,

inner,

strict=True,

force_outplace=False,

return_inputs=False,

return_inputs_states=False,

):

super().__init__()

# inner may be a Module, or it may be an arbitrary callable

# If it's a Module, we get its parameters automatically, which lets

# us avoid a special casing functions versus modules.

self.inner = inner

self.strict = strict

self._force_outplace = force_outplace

self._return_inputs = return_inputs

self._return_inputs_states = return_inputs_states

def forward(self, *args: torch.Tensor):

in_vars, in_desc = _flatten(args)

# NOTE: use full state, because we need it for BatchNorm export

# This differs from the compiler path, which doesn't support it at the moment.

module_state = list(_unique_state_dict(self, keep_vars=True).values())

ret_inputs = []

inputs_states = []

outs = []

def wrapper(*args):

in_args: list[torch.Tensor] = []

for i in range(len(in_vars)):

if not isinstance(args[i], torch.Tensor):

raise RuntimeError("Expected Tensor argument")

in_args.append(args[i])

trace_inputs = _unflatten(in_args, in_desc)

if self._return_inputs:

ret_inputs.append(

tuple(x.clone(memory_format=torch.preserve_format) for x in args)

)

if self._return_inputs_states:

inputs_states.append(_unflatten(in_args, in_desc))

outs.append(self.inner(*trace_inputs))

if self._return_inputs_states:

inputs_states[0] = (inputs_states[0], trace_inputs)

out_vars, _ = _flatten(outs)

if len(out_vars) == 1:

return out_vars[0]

else:

return tuple(out_vars)

graph, _out = torch._C._create_graph_by_tracing(

wrapper,

in_vars + module_state,

_create_interpreter_name_lookup_fn(),

self.strict,

self._force_outplace,

)

if self._return_inputs:

return graph, outs[0], ret_inputs[0]

if self._return_inputs_states:

return graph, outs[0], inputs_states[0]

else:

def clone_input(a):

if a is None:

return None

elif isinstance(a, torch.Tensor):

# TODO: figure out one liner to .clone() and set requires_grad

v = (

a.detach()

.clone(memory_format=None if a.is_mkldnn else torch.preserve_format)

.requires_grad_(a.requires_grad)

)

if a.grad is not None:

v.grad = clone_input(v.grad)

return v

else:

return a.clone(memory_format=torch.preserve_format)

# pyrefly: ignore [missing-attribute]

return function._nested_map(

lambda x: isinstance(x, torch.Tensor), clone_input, condition_msg="tensors"

if (not _JIT_TIME and not time) or not torch.cuda.is_available():

yield

return

stream = torch.cuda.current_stream()

start = torch.cuda.Event(enable_timing=True)

end = torch.cuda.Event(enable_timing=True)

stream.record_event(start)

try:

yield

finally:

stream.record_event(end)

end.synchronize()

"""

# TODO: In principle, we track device information in our trace, so it

# should be possible to check if our execution actually obeyed the 'devices'

# the user provided.

# TODO: Consider adding a utility function to torch.jit to test

# for this case

if not isinstance(model, torch._C.CompiledFunction): # type: ignore[attr-defined]

raise TypeError(

"Cannot verify an uncompiled module. Add @torch.jit.compile to compile it"

)

is_module = isinstance(model, Module)

if not isinstance(args, tuple):

args = (args,)

if is_module:

saved_state = copy.deepcopy(model.state_dict())

def run_fwd_bwd(args, force_trace=False, assert_compiled=False):

params = list(model.parameters()) if is_module else []

in_vars, _ = _flatten((args, params))

# We use a special API to reset the trace and compile it from scratch.

compiled_fn = model

if force_trace:

compiled_fn.clear_cache()

if assert_compiled:

hits = compiled_fn.hits

out = model(*args)

if assert_compiled and compiled_fn.hits == hits: # type: ignore[possibly-undefined]

raise RuntimeError("failed to use the compiled function")

if not isinstance(out, tuple):

out = (out,)

if loss_fn == torch.sum and len(out) != 1:

raise ValueError(

f"Model returns {len(out)} outputs, but default loss function "

"(torch.sum) can only handle a single output"

)

out_vars, _ = _flatten(out)

saved_outs = [

v.detach().clone(memory_format=torch.preserve_format) for v in out_vars

]

loss = loss_fn(*out)

grads = torch.autograd.grad([loss], in_vars)

# TODO: I'm not sure if the clone here is necessary but it is safer

saved_grads = [

v.detach().clone(memory_format=torch.preserve_format) for v in grads

]

return (saved_outs, saved_grads)

with torch.random.fork_rng(devices, _caller="torch.jit.verify"):

uncompiled_outs, uncompiled_grads = run_fwd_bwd(args, force_trace=True)

if not model.has_trace_for(*args):

raise AssertionError("Model should have trace for the given args")

if is_module:

model.load_state_dict(saved_state) # type: ignore[possibly-undefined]

compiled_outs, compiled_grads = run_fwd_bwd(args, assert_compiled=True)

_verify_equal(uncompiled_outs, compiled_outs)

def __init__(self, graph_diff_error, tensor_compare_error, extra_msg=None):

self.message = "Tracing failed sanity checks!\n"

if extra_msg is not None:

self.message += extra_msg + "\n"

if graph_diff_error is not None:

self.message += "ERROR: Graphs differed across invocations!\n"

self.message += indent(graph_diff_error) + "\n"

if tensor_compare_error is not None:

self.message += (

"ERROR: Tensor-valued Constant nodes differed in value "

"across invocations. This often indicates that the tracer has"

" encountered untraceable code.\n"

)

self.message += indent(tensor_compare_error) + "\n"

check_inputs,

func,

traced_func,

check_tolerance,

strict,

force_outplace,

is_trace_module,

_module_class,

example_inputs_is_kwarg=False,

):

# Note: tracing is independent of optimizations, which consume the trace

for inputs in check_inputs:

if isinstance(inputs, torch.Tensor):

inputs = (inputs,)

if is_trace_module:

copied_dict = {}

for name, data in inputs.items():

copied_dict[name] = _clone_inputs(data)

check_mod = torch.jit.trace_module(

getattr(func, "__self__", func),

copied_dict,

check_trace=False,

strict=strict,

_force_outplace=force_outplace,

_module_class=_module_class,

_compilation_unit=torch._C.CompilationUnit(),

example_inputs_is_kwarg=example_inputs_is_kwarg,

_store_inputs=False,

)

check_mod_func = check_mod._c._get_method(traced_func.name)

inputs = inputs[traced_func.name]

if (

isinstance(inputs, (torch.Tensor))

or isinstance(inputs, dict)

and not example_inputs_is_kwarg

):

inputs = (inputs,)

else:

if example_inputs_is_kwarg:

check_mod = torch.jit.trace(

func,

check_trace=False,

strict=strict,

_force_outplace=force_outplace,

_module_class=_module_class,

example_kwarg_inputs=_clone_inputs(inputs),

_store_inputs=False,

)

else:

check_mod = torch.jit.trace(

func,

_clone_inputs(inputs),

check_trace=False,

strict=strict,

_force_outplace=force_outplace,

_module_class=_module_class,

_store_inputs=False,

)

check_mod_func = check_mod

def graph_diagnostic_info():

mod_canonicalized = torch._C._jit_pass_canonicalize(traced_func.graph)

torch._C._jit_pass_inline(mod_canonicalized)

torch._C._jit_pass_erase_shape_information(mod_canonicalized)

mod_str = str(mod_canonicalized)

mod_str = re.sub(r"___torch_mangle_[0-9]+\.", "", mod_str)

check_canonicalized = torch._C._jit_pass_canonicalize(check_mod_func.graph)

torch._C._jit_pass_inline(check_canonicalized)

torch._C._jit_pass_erase_shape_information(check_canonicalized)

check_str = str(check_canonicalized)

check_str = re.sub(r"___torch_mangle_[0-9]+\.", "", check_str)

graph_diff_errors = None

if mod_str != check_str:

import difflib

graph_diff = difflib.ndiff(

mod_str.splitlines(True), check_str.splitlines(True)

)

graph_diff_errors = "Graph diff:\n" + indent("".join(graph_diff)) + "\n"

for n_mod, n_check in zip(

mod_canonicalized.nodes(), check_canonicalized.nodes()

):

if str(n_mod) != str(n_check):

graph_diff_errors += "First diverging operator:\n"

node_diff = difflib.ndiff(

str(n_mod).splitlines(True), str(n_check).splitlines(True)

)

source_printout = (

"Node diff:\n" + indent("".join(node_diff)) + "\n"

)

mod_stack = n_mod.sourceRange()

if mod_stack:

source_printout += (

"Trace source location:\n" + indent(mod_stack) + "\n"

)

check_stack = n_check.sourceRange()

if check_stack:

source_printout += (

"Check source location:\n" + indent(check_stack) + "\n"

)

graph_diff_errors += source_printout

break # For now, only print out the first pair of nodes that diverges

tensor_compare_errors = None

# Check Tensor-valued constant nodes

for n_mod, n_check in zip(

mod_canonicalized.nodes(), check_canonicalized.nodes()

):

if n_mod.kind() != n_check.kind():

break # Graphs have already diverged

if n_mod.kind() == "prim::Constant" and not (

n_mod.mustBeNone() or n_check.mustBeNone()

):

if not n_mod.hasAttribute("value"):

continue

if n_mod.kindOf("value") != "t" or n_check.kindOf("value") != "t":

continue

mod_tensor_val = n_mod.t("value")

check_tensor_val = n_check.t("value")

try:

torch.testing.assert_close(

mod_tensor_val, check_tensor_val, equal_nan=True

)

except (RuntimeError, AssertionError) as e:

if tensor_compare_errors is None:

tensor_compare_errors = ""

tensor_compare_errors += "Node:\n" + indent(str(n_mod)) + "\n"

compare_stack = n_mod.sourceRange()

if compare_stack:

tensor_compare_errors += (

"Source Location:\n" + indent(compare_stack) + "\n"

)

tensor_compare_errors += "Comparison exception: " + indent(

str(e)

)

break # For now, only print the first diverging pair

return graph_diff_errors, tensor_compare_errors

def wrap_retval(x):

return x if isinstance(x, tuple) else (x,)

def run_mod_and_filter_tensor_outputs(mod, inputs, running_what):

try:

if isinstance(inputs, dict) and example_inputs_is_kwarg:

outs = wrap_retval(mod(**inputs))

else:

outs = wrap_retval(mod(*_clone_inputs(inputs)))

outs = [out for out in outs if isinstance(out, torch.Tensor)]

return outs

except Exception as e:

graph_diff_errors, tensor_compare_errors = graph_diagnostic_info()

msg = f"encountered an exception while running the {running_what} with test inputs.\nException:\n{indent(str(e))}"

raise TracingCheckError(

graph_diff_errors,

tensor_compare_errors,

extra_msg=msg,

) from e

has_warned = [False]

def maybe_warn_nondeterministic():

if has_warned[0]:

return

has_warned[0] = True

nondeterm_ops = [

op for op in traced_func.graph.nodes() if op.isNondeterministic()

]

if len(nondeterm_ops) > 0:

nondeterministic_ops_warning = "Trace had nondeterministic nodes. "

nondeterministic_ops_warning += (

"Did you forget call .eval() on your model? Nodes:\n"

)

nondeterministic_ops_warning += "\n".join(

[indent(str(op)) for op in nondeterm_ops][:20]

)

nondeterministic_ops_warning += (

"\nThis may cause errors in trace checking. To disable trace checking,"

" pass check_trace=False to torch.jit.trace()"

)

warnings.warn(

nondeterministic_ops_warning, category=TracerWarning, stacklevel=5

)

def compare_outputs(original, reference, match_what):

all_ok = True

for i, (orig, ref) in enumerate(zip(original, reference)):

try:

if orig.is_quantized:

orig = orig.dequantize()

if ref.is_quantized:

ref = ref.dequantize()

if orig.is_mkldnn:

orig = orig.to_dense()

if ref.is_mkldnn:

ref = ref.to_dense()

if ref.is_complex() or orig.is_complex():

torch.testing.assert_close(

orig.to(torch.cdouble),

ref.to(torch.cdouble),

rtol=check_tolerance,

atol=default_tolerances(orig, ref)[1],

equal_nan=True,

)

else:

if orig.is_mps or ref.is_mps:

torch.testing.assert_close(

orig.float(),

ref.float(),

rtol=check_tolerance,

atol=default_tolerances(orig, ref)[1],

equal_nan=True,

)

elif getattr(orig, "is_nested", None) or getattr(

ref, "is_nested", None

):

if getattr(orig, "is_nested", None) != getattr(

ref, "is_nested", None

):

raise AssertionError(

f"Nested tensor mismatch: orig.is_nested="

f"{getattr(orig, 'is_nested', None)}, "

f"ref.is_nested={getattr(ref, 'is_nested', None)}"

)

for t_orig, t_ref in zip(orig.unbind(), ref.unbind()):

torch.testing.assert_close(

t_orig.double(),

t_ref.double(),

rtol=check_tolerance,

atol=default_tolerances(t_orig, t_ref)[1],

equal_nan=True,

)

else:

torch.testing.assert_close(

orig.double(),

ref.double(),

rtol=check_tolerance,

atol=default_tolerances(orig, ref)[1],

equal_nan=True,

)

except AssertionError as e:

maybe_warn_nondeterministic()

warnings.warn(

"Output nr "

+ str(i + 1)

+ ". of the traced function does not match "

"the corresponding output of the "

+ match_what

+ ". Detailed error:\n"

+ str(e),

category=TracerWarning,

stacklevel=4,

)

all_ok = False

return all_ok

traced_outs = run_mod_and_filter_tensor_outputs(traced_func, inputs, "trace")

fn_outs = run_mod_and_filter_tensor_outputs(func, inputs, "Python function")

if compare_outputs(traced_outs, fn_outs, "Python function"):

check_outs = run_mod_and_filter_tensor_outputs(

check_mod_func, inputs, "repeated trace"

)

compare_outputs(traced_outs, check_outs, "repeated trace")

diag_info = graph_diagnostic_info()

if any(info is not None for info in diag_info):

@staticmethod

def ignore_lib_warnings():

# We ignore warnings from all submodules excluding the JIT, because we need them e.g. for _check_trace

warnings.filterwarnings(

"ignore", category=TracerWarning, module="torch.(?!jit)"

)

if isinstance(example_inputs, (torch.Tensor, dict)):

return (example_inputs,)

# done primarily so that weird iterables fail here and not pybind11 code

if not isinstance(example_inputs, tuple):

return tuple(example_inputs)

if isinstance(mod, ScriptModule):

return mod

elif torch._jit_internal.module_has_exports(mod):

infer_methods_stubs_fn = torch.jit._recursive.make_stubs_from_exported_methods

return torch.jit._recursive.create_script_module(

mod, infer_methods_stubs_fn, share_types=False, is_tracing=True

)

else:

if _module_class is None:

_module_class = TopLevelTracedModule

if check_inputs is None:

return None

import torch.utils._pytree as pytree

flat_export = pytree.tree_leaves(export)

flat_trace = pytree.tree_leaves(trace)

for orig, loaded in zip(flat_export, flat_trace):

if orig.layout != loaded.layout:

return False

# mkldnn is not supported for torch.allclose

if orig.layout == torch._mkldnn: # type: ignore[attr-defined]

return True

if type(orig) is not type(loaded):

return False

if isinstance(orig, torch._subclasses.FakeTensor):

# Skip for FakeTensor.

return True

elif isinstance(orig, torch.Tensor):

if orig.dtype != loaded.dtype:

return False

if not torch.allclose(orig, loaded):

return False

else:

if orig != loaded:

return False

func,

example_inputs=None,

optimize=None,

check_trace=True,

check_inputs=None,

check_tolerance=1e-5,

strict=True,

_force_outplace=False,

_module_class=None,

_compilation_unit=_python_cu,

example_kwarg_inputs=None,

_store_inputs=True,

):

if isinstance(func, torch.jit.ScriptModule):

# it is hard to trace it because the forward method on ScriptModule is already defined, so it

# would result in an error.

warnings.warn(

"The input to trace is already a ScriptModule, tracing it is a no-op. Returning the object as is.",

stacklevel=2,

)

return func

if isinstance(func, torch.nn.Module):

if example_inputs is None:

if isinstance(example_kwarg_inputs, dict):

example_inputs = example_kwarg_inputs

else:

raise RuntimeError("example_kwarg_inputs should be a dict")

return trace_module(

func,

{"forward": example_inputs},

None,

check_trace,

wrap_check_inputs(check_inputs),

check_tolerance,

strict,

_force_outplace,

_module_class,

example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),

_store_inputs=_store_inputs,

)

if (

hasattr(func, "__self__")

and isinstance(func.__self__, torch.nn.Module)

and func.__name__ == "forward"

):

if example_inputs is None:

if isinstance(example_kwarg_inputs, dict):

example_inputs = example_kwarg_inputs

else:

raise RuntimeError("example_kwarg_inputs should be a dict")

return trace_module(

func.__self__,

{"forward": example_inputs},

None,

check_trace,

wrap_check_inputs(check_inputs),

check_tolerance,

strict,

_force_outplace,

_module_class,

example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),

_store_inputs=_store_inputs,

)

# Special case for common case of passing a single Tensor

if (

isinstance(example_inputs, (torch.Tensor, dict))

and example_kwarg_inputs is None

):

example_inputs = (example_inputs,)

# done primarily so that weird iterables fail here and not pybind11 code

elif example_kwarg_inputs is None and not isinstance(example_inputs, tuple):

# pyrefly: ignore [bad-argument-type]

example_inputs = tuple(example_inputs)

var_lookup_fn = _create_interpreter_name_lookup_fn(0)

if hasattr(func, "__self__") and isinstance(func.__self__, torch.nn.Module):

raise AttributeError(

"trace doesn't support compiling individual module's functions.\n"

"Please use trace_module"

)

name = _qualified_name(func)

if isinstance(example_kwarg_inputs, dict):

example_inputs = example_kwarg_inputs

traced = torch._C._create_function_from_trace_with_dict(

name,

func,

example_kwarg_inputs,

var_lookup_fn,

strict,

_force_outplace,

get_callable_argument_names(func),

)

else:

traced = torch._C._create_function_from_trace(

name,

func,

# pyrefly: ignore [bad-argument-type]

example_inputs,

var_lookup_fn,

strict,

_force_outplace,

get_callable_argument_names(func),

)

# Check the trace against new traces created from user-specified inputs

if check_trace:

if check_inputs is not None:

_check_trace(

check_inputs,

func,

traced,

check_tolerance,

strict,

_force_outplace,

False,

_module_class,

example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),

)

else:

_check_trace(

[example_inputs],

func,

traced,

check_tolerance,

strict,

_force_outplace,

False,

_module_class,

example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),

)

# Allow torch.compile() to inline

traced._torchdynamo_inline = func # type: ignore[attr-defined]

DIRECT_EXPORT = "DIRECT_EXPORT"

TRACE_AND_EXPORT = "TRACE_AND_EXPORT"

SOURCE_TO_SOURCE = "SOURCE_TO_SOURCE"

def __str__(self) -> str:

SUCCESS = "SUCCESS"

FAILED_TO_EXPORT = "FAILED_TO_EXPORT"

FAILED_TO_RUN = "FAILED_TO_RUN"

ACCURACY_ERROR = "ACCURACY_ERROR"

def __str__(self) -> str:

func,

example_inputs=None,

optimize=None,

check_trace=True,

check_inputs=None,

check_tolerance=1e-5,

strict=True,

_force_outplace=False,

_module_class=None,

_compilation_unit=_python_cu,

example_kwarg_inputs=None,

_store_inputs=True,

):

r"""

if sys.version_info >= (3, 14):

warnings.warn(