[pytree] Extract reusable generic tests for pytree (#110395)

Part of #109684

- #109684

Changes:

- Add new functions `tree_structure`, `tree_leaves`, `tree_map_` and `tree_map_only_` to Python pytree.
- Extract reusable tests for pytree to `TestGenericPytree`.
- Change `treespec_dumps` and `treespec_loads` in C++ pytree to call Python pytree and use JSON string as serialization type.
- Rename `torch.utils.pytree` -> `torch.utils._cxx_pytree`.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/110395
Approved by: https://github.com/zou3519
This commit is contained in:
Xuehai Pan
2023-10-04 23:40:43 +00:00
committed by PyTorch MergeBot
parent 37afa0c349
commit 449271f3f1
6 changed files with 428 additions and 416 deletions

View File

@ -43,7 +43,7 @@ files =
tools,
torch/profiler/_memory_profiler.py,
torch/utils/_pytree.py,
torch/utils/pytree.py,
torch/utils/_cxx_pytree.py,
torch/utils/benchmark/utils/common.py,
torch/utils/benchmark/utils/timer.py,
torch/utils/benchmark/utils/valgrind_wrapper