Animesh Jain 4308b8a28f [dynamo] Support torch.fx.traceback.annotate (#164678) ...

Builds on top of https://github.com/pytorch/pytorch/pull/163673 and https://github.com/pytorch/pytorch/pull/164174. This will be used in the followup PRs to apply regional inductor compilation.

The existing implementation let Dynamo trace into the `torch.fx.traceback.annotate`, but thats not what we want. We want Dynamo to essentially run the torch.fx.traceback.annotate function in eager, so that every Fx node created in Dynamo Fx graph has the custom meta node.

What does not work?
* We still have to set the context manager `torch.fx.traceback.preserve_node_meta()` in the user code because CI was unhappy. This can be fixed but with some perseverance.
* This does not work with graph breaks yet. But we can solve that problem, if needed, in a separate PR.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/164678
Approved by: https://github.com/SherlockNoMad, https://github.com/jansel, https://github.com/xmfan

2025-10-08 22:41:00 +00:00

..

__init__.py

[export] Fix test to run internally (#116118)

2023-12-20 01:02:16 +00:00

opinfo_schema.py

Revert "Enable all SIM rules except disabled ones (#164645)"

2025-10-05 19:32:21 +00:00

random_dag.py

[BE] fix ruff rule E226: add missing whitespace around operator in f-strings (#144415)

2025-01-08 21:55:00 +00:00

test_converter.py

Enable more nightly tests on s390x (#160893)

2025-08-28 22:20:55 +00:00

test_cpp_serdes.py

[export] Implement cpp deserializer. (#136398)

2024-11-14 16:34:59 +00:00

test_db.py

Move export_db to use new tracer, remove restriction on optional inputs (#162993)

2025-09-18 00:43:32 +00:00

test_draft_export.py

[dynamic shapes] unbacked-safe slicing (#161414)

2025-09-30 01:15:19 +00:00

test_dynamic_shapes.py

[Dynamic Shape][BE] trim _DimHint serialization (#163891)

2025-09-27 00:08:01 +00:00

test_experimental.py

Revert "Fix double dispatch to Python for detach (#163671)"

2025-10-08 01:46:45 +00:00

test_export_opinfo.py

Enable torch.nn.functional.batch_norm in test_export_opinfo (#164261)

2025-10-01 21:56:08 +00:00

test_export_strict.py

Set non-strict export as default mode (#148790)

2025-03-12 21:10:58 +00:00

test_export_training_ir_to_run_decomp.py

Replace export_for_training with export (#162396)

2025-09-10 14:19:34 +00:00

test_export_with_inline_and_install.py

[dynamo][export] Add some missing trace rules (#164080)

2025-09-29 08:47:24 +00:00

test_export.py

[dynamo] Support torch.fx.traceback.annotate (#164678)

2025-10-08 22:41:00 +00:00

test_functionalized_assertions.py

fix incorrect literal strings / accidental tuples (#146037)

2025-02-03 15:08:11 +00:00

test_hop.py

Remove unused import in tests (#146266)

2025-02-03 03:40:18 +00:00

test_lift_unlift.py

Consistently use load_torchbind_test_lib in tests (#148082)

2025-03-03 19:37:28 +00:00

test_nativert.py

Stop building nativert in OSS (#164463)

2025-10-03 19:41:15 +00:00

test_package.py

[export] Exporter API prototype. (#153205)

2025-05-11 14:20:09 +00:00

test_pass_infra.py

[Codemod][AddExplicitStrictExportArg] caffe2/test (#143688)

2024-12-27 07:58:44 +00:00

test_passes.py

Revert "Enable all SIM rules except disabled ones (#164645)"

2025-10-05 19:32:21 +00:00

test_retraceability.py

Set non-strict export as default mode (#148790)

2025-03-12 21:10:58 +00:00

test_schema.py

[resubmit][export] add _union_dataclass to support comparing dataclasses that inherits from union. (#156765)

2025-06-25 22:32:12 +00:00

test_serdes.py

Set non-strict export as default mode (#148790)

2025-03-12 21:10:58 +00:00

test_serialize.py

use check_size instead of check_is_size in ops.py (#164668)

2025-10-08 14:23:38 +00:00

test_sparse.py

Replace all random is_fbcode imports to environment (#151283)

2025-04-15 19:42:58 +00:00

test_strict_export_v2.py

Fix inconsistent test and add new tracer as config (#162558)

2025-09-17 17:01:48 +00:00

test_swap.py

Fix inconsistent test and add new tracer as config (#162558)

2025-09-17 17:01:48 +00:00

test_tools.py

Add None return type to init -- tests (#132352)

2024-08-01 15:44:51 +00:00

test_torchbind.py

[RELAND v2] Close some sources of fake tensors (#164372)

2025-10-02 18:58:52 +00:00

test_tree_utils.py

fix incorrect literal strings / accidental tuples (#146037)

2025-02-03 15:08:11 +00:00

test_unflatten_training_ir.py

Replace export_for_training with export (#162396)

2025-09-10 14:19:34 +00:00

test_unflatten.py

shape guards (#161178)

2025-09-08 22:44:09 +00:00

test_upgrader.py

[schema_upgrader] add C++ upgrader for json based upgrading (#156761)

2025-06-28 18:15:06 +00:00

test_verifier.py

Replace export_for_training with export (#162396)

2025-09-10 14:19:34 +00:00

testing.py

testing infra and some fixes (#162183)

2025-09-10 20:48:12 +00:00