Revert "[1/N] Remove unused loop variables (pytorch#166258)"

This reverts commit 76b2c37.

Reverted pytorch#166258 on behalf of https://github.com/atalman due to breaks test/distributed/test_serialization.py::TestSerialization::test_weights_only [GH job link](https://github.com/pytorch/pytorch/actions/runs/18894311802/job/53929321703) [HUD commit link](https://hud.pytorch.org/pytorch/pytorch/commit/76b2c37045e52540ec51e967aa7b6436a6b9b174) ([comment](pytorch#166258 (comment)))

Author: pytorchmergebot

torch/_dynamo/eval_frame.py

@@ -1791,7 +1791,7 @@ def produce_matching(
         for i, val in enumerate(sources):
             dict_of_source_vals[id(val)] = i
 
-        for val in candidates:
+        for i, val in enumerate(candidates):
             if isinstance(val, tuple(common_constant_types)):
                 matched_elements_positions.append(None)
             elif id(val) not in dict_of_source_vals:

torch/_dynamo/guards.py

@@ -317,7 +317,7 @@ def visit_dict_manager(node: DictGuardManager) -> bool:
             is_diff_guard_node = (
                 node.get_source() in self.diff_guard_sources or node.fail_count() > 0
             )
-            for _idx, (key_mgr, val_mgr) in sorted(
+            for idx, (key_mgr, val_mgr) in sorted(
                 node.get_key_value_managers().items()
             ):
                 is_diff_guard_node |= visit(key_mgr) | visit(val_mgr)
@@ -440,15 +440,17 @@ def visit_dict_manager(node: DictGuardManager) -> list[GuardManager]:
             is_subtree_tag_safe = True
 
             # Recurse to get the tag safe roots from subtree.
-            for _idx, (key_mgr, val_mgr) in sorted(
+            for idx, (key_mgr, val_mgr) in sorted(
                 node.get_key_value_managers().items()
             ):
                 if key_mgr is not None:
                     visit(key_mgr)
                 if val_mgr is not None:
                     tag_safe_roots.extend(visit(val_mgr))
 
-            for key_mgr, val_mgr in node.get_key_value_managers().values():
+            for idx, (key_mgr, val_mgr) in sorted(
+                node.get_key_value_managers().items()
+            ):
                 if key_mgr:
                     is_subtree_tag_safe &= key_mgr.is_tag_safe()
 

torch/_dynamo/variables/optimizer.py

@@ -289,7 +289,9 @@ def mark_static(x):
             params_vt = group_vt.getitem_const(tx, ConstantVariable.create("params"))
             all_static = True
             non_static_grads = []
-            for p, p_vt in zip(group["params"], params_vt.unpack_var_sequence(tx)):
+            for p_ind, (p, p_vt) in enumerate(
+                zip(group["params"], params_vt.unpack_var_sequence(tx))
+            ):
                 param_source = p_vt.source
                 self.tensor_to_source[p] = param_source
                 grad_source = GradSource(
@@ -320,12 +322,12 @@ def mark_static(x):
 
         # We have to again iterate over the state dict to collect the
         # tensor_to_source dict. This is used for the finalizer.
-        for idx, value in enumerate(self.value.state.values()):
+        for idx, (p, value) in enumerate(self.value.state.items()):
             p_state_source = DictGetItemSource(
                 state_source, ConstDictKeySource(state_source, idx)
             )
             tx.output.guard_on_key_order.add(p_state_source)
-            for inner_idx, v in enumerate(value.values()):
+            for inner_idx, (k, v) in enumerate(value.items()):
                 if (
                     isinstance(v, torch.Tensor)
                     and v not in self.grad_to_source

torch/_functorch/_aot_autograd/collect_metadata_analysis.py

@@ -240,7 +240,7 @@ def inner(*flat_args):
 
         # Inspect the state of the input tensor functional wrapper to detect input mutation info
         # If inp[i] has a metadata-only mutation, then maybe_inputs_with_mutated_metadata[i] contains the updated version
-        for arg, f_arg in zip(flat_args, flat_f_args):
+        for i, (arg, f_arg) in enumerate(zip(flat_args, flat_f_args)):
             # NB: Mutation of non-contiguous tensor subclass input can result in a mismatch in
             # strides between the functionalized arg inner tensors and non-functionalized arg inner
             # tensors. This is a problem as the inner tensor stride change may not be reflected

torch/_functorch/_aot_autograd/runtime_wrappers.py

@@ -2041,7 +2041,7 @@ def maybe_coerce(x):
 
         assert len(meta.attrs) == len(runtime_subclass_keys)
         leaves = []
-        for attr, attr_meta in meta.attrs.items():
+        for i, (attr, attr_meta) in enumerate(meta.attrs.items()):
             elem = getattr(x, attr)
             new_elem, elem_leaves = AOTDispatchAutograd.process_runtime_tangent(
                 elem, attr_meta

torch/_functorch/_aot_autograd/subclass_parametrization.py

@@ -98,7 +98,7 @@ def unwrap_tensor_subclass_parameters(module: torch.nn.Module) -> torch.nn.Modul
                 module, name, UnwrapTensorSubclass()
             )
 
-    for child in module.children():
+    for name, child in module.named_children():
         unwrap_tensor_subclass_parameters(child)
 
     return module

torch/_functorch/partitioners.py

@@ -1481,7 +1481,9 @@ def get_sample_rng_state(device: Optional[torch.device]):
         )
     )
 
-    for rng_count, node_pair in enumerate(recomputable_rng_ops_map.values()):
+    for rng_count, (base_node, node_pair) in enumerate(
+        recomputable_rng_ops_map.items()
+    ):
         # Step 2 - Modify the fwd pass such that
         fw_node = node_pair["fwd"]
         bw_node = node_pair["bwd"]
@@ -2712,7 +2714,9 @@ def thread_graphsafe_rng_from_hops(module, is_backward):
         subgraph = getattr(module, hop_node.args[0].target)
         if isinstance(subgraph, fx.GraphModule):
             new_rng_inputs = []
-            for placeholder_node in subgraph.graph.find_nodes(op="placeholder"):
+            for idx, placeholder_node in enumerate(
+                subgraph.graph.find_nodes(op="placeholder")
+            ):
                 if rng_string in placeholder_node.name:
                     # Found a rng state placeholder in the hop graph, lets add
                     # the corresponding node in the outer graph

torch/_functorch/pyfunctorch.py

@@ -116,7 +116,7 @@ def temporarily_restore_interpreter_stack(stack):
             pushed.append(s)
         yield
     finally:
-        for _ in reversed(pushed):
+        for s in reversed(pushed):
             # TODO: would be nice to assert that the layers are the same, but
             # Python object identity is not preserved
             pop_dynamic_layer_stack()

torch/_higher_order_ops/utils.py

@@ -907,7 +907,7 @@ def diff_tensor_meta(
         try:
             if val1 != val2:
                 pair_diffs.append(f"'{meta_name}: {val1} vs {val2}'")
-        except GuardOnDataDependentSymNode:
+        except GuardOnDataDependentSymNode as _:
             pair_diffs.append(f"'{meta_name}: {val1} vs {val2}'")
             continue
     return pair_diffs
@@ -1197,7 +1197,7 @@ def wrapped_fn(*flat_args):
 
     # call_op preserves ordering of proxies via schema
     materialized_args = []
-    for i, proxy in enumerate(arg_proxies):
+    for i, (proxy, arg) in enumerate(zip(arg_proxies, schema.arguments)):
         if (
             isinstance(proxy, torch.fx.Node)
             and proxy.op == "get_attr"

torch/_higher_order_ops/while_loop.py

@@ -316,7 +316,7 @@ def _validate_cond_output(pred):
 
     if stack_output:
         outs: list[torch.Tensor] = []
-        for out in outputs:
+        for i, out in enumerate(outputs):
             outs.append(torch.stack(out, dim=0))
         return tuple(outs)