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

Helps with reducing Dynamo tracing time. Earlier the generator object
would cause more polyfills.

Pull Request resolved: pytorch#168024
Approved by: https://github.com/williamwen42

Author: anijain2305

torch/_dynamo/polyfills/pytree.py

@@ -201,8 +201,11 @@ def __post_init__(self, /) -> None:
             num_children = 0
         else:
             assert callable(self._unflatten_func)
-            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_nodes = 1
+            num_leaves = 0
+            for child in self._children:
+                num_nodes += child.num_nodes
+                num_leaves += child.num_leaves
             num_children = len(self._children)
 
         object.__setattr__(self, "num_nodes", num_nodes)

torch/utils/_pytree.py

@@ -1113,8 +1113,11 @@ def __post_init__(self) -> None:
             num_leaves = 1
             num_children = 0
         else:
-            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_nodes = 1
+            num_leaves = 0
+            for child in self._children:
+                num_nodes += child.num_nodes
+                num_leaves += child.num_leaves
             num_children = len(self._children)
         object.__setattr__(self, "num_nodes", num_nodes)
         object.__setattr__(self, "num_leaves", num_leaves)