pytorch/torch/utils
Henry Hu 908133f682 [TreeSpec] Add custom comparision function (#146442)
Summary:
https://github.com/pytorch/pytorch/pull/145815 used caching to for treespec_loads calculation to speed up AOTI module call.

However, this made tests flaky due when comparing TreeSpec for objects in local scope. ie. 'test_export.TestExport.test_pytree_register_nested_data_class.<locals>.Inner'

Type comparison will yield False when local scopes are different due to lru_cache.

Since this comparison is only used for testing purpose, we will only test if str(type) are equal.

Test Plan:
```
PYTORCH_TEST_WITH_ROCM=1 python test/export/test_retraceability.py
```

Differential Revision: D69137706

Pull Request resolved: https://github.com/pytorch/pytorch/pull/146442
Approved by: https://github.com/angelayi
2025-02-07 22:39:21 +00:00
..
_strobelight PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_sympy Revert "[inductor] Refactor op handlers part 5 (#146257)" 2025-02-05 23:20:38 +00:00
backcompat
benchmark PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
bottleneck PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
data Fixed a typo in dataset.py (#146600) 2025-02-07 05:09:51 +00:00
hipify PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
jit PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
model_dump PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
serialization Add option to serialization config to reduce random reads from get_record_offset when loading with mmap=True (#143880) 2025-01-31 17:09:20 +00:00
tensorboard PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
viz PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
__init__.py
_backport_slots.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_config_module.py Add check that envvar configs are boolean (#145454) 2025-02-05 19:40:10 +00:00
_config_typing.pyi
_content_store.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_contextlib.py
_cpp_embed_headers.py PEP585: Missed conversions (#145342) 2025-01-29 05:24:36 +00:00
_cpp_extension_versioner.py
_cxx_pytree.py [AOTI] Cache treespec_loads calculation (#145815) 2025-01-31 06:38:21 +00:00
_device.py Revert "Fix torch.normal ignores default_device (#144070)" 2025-01-14 17:41:58 +00:00
_exposed_in.py Type exposed_in decorator (#141894) 2024-12-03 16:28:17 +00:00
_filelock.py filelock: Make waitcounter variant to use (#139816) 2024-12-12 01:18:34 +00:00
_foreach_utils.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_freeze.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_functools.py dynamo tracing perf: cache on import_source: 52.9 -> 52.58 (#143058) 2024-12-13 18:20:48 +00:00
_get_clean_triton.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_import_utils.py remove allow-untyped-defs from utils/_import_utils.py (#144089) 2025-01-03 01:21:13 +00:00
_mode_utils.py
_ordered_set.py [BE]: Update typing of OrderedSet ancestor (#145783) 2025-01-28 04:43:49 +00:00
_python_dispatch.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_pytree.py [TreeSpec] Add custom comparision function (#146442) 2025-02-07 22:39:21 +00:00
_stats.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_thunk.py
_traceback.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
_triton.py Adding lowering to persistent-tma device kernel for _scaled_mm (#142045) 2024-12-09 01:48:40 +00:00
_typing_utils.py Revert "Fix type annotation of Linear.bias (#142326)" 2025-01-26 03:41:00 +00:00
_zip.py
backend_registration.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
bundled_inputs.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
checkpoint.py [BE] typing for decorators (#144161) 2025-01-04 16:40:09 +00:00
collect_env.py Revert "Collect packages with importlib in collect_env (#144616)" 2025-01-13 03:11:04 +00:00
cpp_backtrace.py
cpp_extension.py Set -DPy_LIMITED_API flag for py_limited_api=True extensions (#145764) 2025-01-28 20:11:05 +00:00
deterministic.py
dlpack.py
file_baton.py
flop_counter.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
hooks.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
mkldnn.py
mobile_optimizer.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
model_zoo.py
module_tracker.py PEP585 update - torch/utils (#145201) 2025-01-21 21:04:10 +00:00
show_pickle.py Use typing.IO[bytes] instead of io.BytesIO in annotations (#144994) 2025-01-27 18:08:07 +00:00
throughput_benchmark.py
weak.py