Revert "[pytree][compile] Slightly faster TreeSpec init (pytorch#168024)"

This reverts commit db1551b.

Reverted pytorch#168024 on behalf of https://github.com/yangw-dev due to Internal merge fail, These changes have conflicts when merging with master branch. Rebase this diff. please rebase the pr and try merge again ([comment](pytorch#168024 (comment)))

Author: pytorchmergebot

torch/_dynamo/polyfills/pytree.py

@@ -201,11 +201,8 @@ def __post_init__(self, /) -> None:
             num_children = 0
         else:
             assert callable(self._unflatten_func)
-            num_nodes = 1
-            num_leaves = 0
-            for child in self._children:
-                num_nodes += child.num_nodes
-                num_leaves += child.num_leaves
+            num_nodes = sum((spec.num_nodes for spec in self._children), start=1)
+            num_leaves = sum(spec.num_leaves for spec in self._children)
             num_children = len(self._children)
 
         object.__setattr__(self, "num_nodes", num_nodes)

torch/utils/_pytree.py

@@ -1113,11 +1113,8 @@ def __post_init__(self) -> None:
             num_leaves = 1
             num_children = 0
         else:
-            num_nodes = 1
-            num_leaves = 0
-            for child in self._children:
-                num_nodes += child.num_nodes
-                num_leaves += child.num_leaves
+            num_nodes = sum((spec.num_nodes for spec in self._children), start=1)
+            num_leaves = sum(spec.num_leaves for spec in self._children)
             num_children = len(self._children)
         object.__setattr__(self, "num_nodes", num_nodes)
         object.__setattr__(self, "num_leaves", num_leaves)