pytorch/torch/utils
Xuehai Pan 108bb224f7 [pytree] add treespec_{leaf,tuple,dict} functions for args_spec modification (#160843)
The goal of this PR is to provide a standard way to create simple treespec instances and hide the implementation details of the `PyTreeSpec` class.

Changes:

1. Add function `treespec_leaf()` to replace `LeafSpec()`.
2. Add function `treespec_tuple(...)` and `treespec_dict(...)` to create treespec for `tuple` / `dict` which is used for `*args` / `**kwargs`. This avoids direct modification to `treespec` instances that rely on the implementation details of the `PyTreeSpec` class.
3. Change `len(spec.children_specs)` to `spec.num_children`.
4. Change `isinstance(spec, LeafSpec)` to `spec.is_leaf()`.

------

Pull Request resolved: https://github.com/pytorch/pytorch/pull/160843
Approved by: https://github.com/mlazos
2025-10-31 10:33:16 +00:00
..
_strobelight Fix error suppression syntax in utils and nn (#166242) 2025-10-26 05:21:07 +00:00
_sympy Revert "[2/N] Add strict parameter to Python zip calls (#166257)" 2025-10-30 13:20:00 +00:00
backcompat
benchmark Revert "[2/N] Add strict parameter to Python zip calls (#166257)" 2025-10-30 13:20:00 +00:00
data Revert "[2/N] Add strict parameter to Python zip calls (#166257)" 2025-10-30 13:20:00 +00:00
hipify Fix pyrefly ignores 1/n (#166239) 2025-10-26 00:44:10 +00:00
jit [torch/utils][Code Clean] Clean asserts in hipify/, jit/, model_dump and tensorboard of torch/utils (#165311) 2025-10-14 15:26:23 +00:00
model_dump Fix flake8 B028 warnings (#166224) 2025-10-26 06:18:55 +00:00
serialization
tensorboard Revert "[2/N] Add strict parameter to Python zip calls (#166257)" 2025-10-30 13:20:00 +00:00
viz Revert "[2/N] Add strict parameter to Python zip calls (#166257)" 2025-10-30 13:20:00 +00:00
__init__.py
_appending_byte_serializer.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
_config_module.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
_config_typing.pyi [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
_content_store.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
_contextlib.py added type annotation to _NoParamDecoratorContextManager.__new__ (#166414) 2025-10-28 21:59:20 +00:00
_cpp_embed_headers.py [Fix] Adding missing f prefixes to formatted strings [2/N] (#164066) 2025-09-29 04:40:44 +00:00
_cpp_extension_versioner.py
_cxx_pytree.py [pytree] add treespec_{leaf,tuple,dict} functions for args_spec modification (#160843) 2025-10-31 10:33:16 +00:00
_debug_mode.py [DebugMode] store stringify args by default (#166347) 2025-10-30 22:12:23 +00:00
_device.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
_dtype_abbrs.py
_exposed_in.py [2/N] Apply ruff UP035 check in torch files (#164054) 2025-09-29 03:35:32 +00:00
_filelock.py
_foreach_utils.py [2/N] Apply ruff UP035 check in torch files (#164054) 2025-09-29 03:35:32 +00:00
_functools.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
_get_clean_triton.py [BE][Ez]: Use sys.executable instead of hardcoded Python (#165679) 2025-10-17 01:07:40 +00:00
_helion.py
_import_utils.py
_mode_utils.py
_ordered_set.py Fix pyrefly ignores 1/n (#166239) 2025-10-26 00:44:10 +00:00
_python_dispatch.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
_pytree.py [pytree] add treespec_{leaf,tuple,dict} functions for args_spec modification (#160843) 2025-10-31 10:33:16 +00:00
_stats.py [2/N] Apply ruff UP035 check in torch files (#164054) 2025-09-29 03:35:32 +00:00
_thunk.py [2/N] Apply ruff UP035 check in torch files (#164054) 2025-09-29 03:35:32 +00:00
_traceback.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
_triton.py Remove unnecessary noqa suppressions (#164106) 2025-10-18 04:52:41 +00:00
_typing_utils.py
_zip.py
backend_registration.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
bundled_inputs.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
checkpoint.py Use 'is' in callable comparisons (#166624) 2025-10-30 19:00:09 +00:00
collect_env.py [2/N] Fix ruff warnings (#164460) 2025-10-04 03:40:32 +00:00
cpp_backtrace.py
cpp_extension.py Revert "[2/N] Add strict parameter to Python zip calls (#166257)" 2025-10-30 13:20:00 +00:00
deterministic.py
dlpack.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
file_baton.py Fix flake8 B028 warnings (#166224) 2025-10-26 06:18:55 +00:00
flop_counter.py Revert "[2/N] Add strict parameter to Python zip calls (#166257)" 2025-10-30 13:20:00 +00:00
hooks.py Revert "[2/N] Add strict parameter to Python zip calls (#166257)" 2025-10-30 13:20:00 +00:00
mkldnn.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
mobile_optimizer.py
model_zoo.py
module_tracker.py
show_pickle.py
throughput_benchmark.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00
weak.py [torch/utils][Code Clean] Clean asserts in torch/utils/*.py (#165410) 2025-10-27 19:48:55 +00:00