pytorch/torch/_export/serde
Xuehai Pan e8fadba28c [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-11-01 04:12:11 +00:00
..
gen-cpp2 [export][ez] Fix forward D67044185 (#143193) 2024-12-13 16:06:42 +00:00
__init__.py
dynamic_shapes.py Fix pyrefly ignore syntax (#166438) 2025-10-29 00:02:21 +00:00
export_schema.thrift shape guards (#161178) 2025-09-08 22:44:09 +00:00
schema_check.py [2/N] Use 'is' in callable comparisons (#166685) 2025-10-31 08:08:07 +00:00
schema.py docs: fix typos (#164879) 2025-10-28 12:00:36 +00:00
schema.yaml shape guards (#161178) 2025-09-08 22:44:09 +00:00
serialize.py [pytree] add treespec_{leaf,tuple,dict} functions for args_spec modification (#160843) 2025-11-01 04:12:11 +00:00
union.py [BE][2/16] fix typos in torch/ (torch/_*/) (#156312) 2025-07-12 05:47:06 +00:00