cast to node to appease linter

danielvegamyhre · danielvegamyhre · commit 593119fec28c · 2025-03-20T05:31:55.000-07:00
diff --git a/torch/_inductor/fx_passes/micro_pipeline_tp.py b/torch/_inductor/fx_passes/micro_pipeline_tp.py
@@ -655,7 +655,9 @@ def _scatter_dim_after_reshape(
         "reshape must produce 2D tensor for scaled_mm"
     )
 
-    reshape_op_input_tensor = _get_tensor(reshape_node.args[0])
+    assert len(reshape_node.args) >= 1, "reshape node must have at least 1 arg"
+    input_tensor_node = cast(torch.fx.Node, reshape_node.args[0])
+    reshape_op_input_tensor = _get_tensor(input_tensor_node)
     assert reshape_op_input_tensor.ndim > reshape_op_output_tensor.ndim, (
         "reshape must be from 3D+ to 2D"
     )

Original file line number	Diff line number	Diff line change
`@@ -655,7 +655,9 @@ def _scatter_dim_after_reshape(`
`655`	`655`	`"reshape must produce 2D tensor for scaled_mm"`
`656`	`656`	`)`
`657`	`657`
`658`		`- reshape_op_input_tensor = _get_tensor(reshape_node.args[0])`
	`658`	`+ assert len(reshape_node.args) >= 1, "reshape node must have at least 1 arg"`
	`659`	`+ input_tensor_node = cast(torch.fx.Node, reshape_node.args[0])`
	`660`	`+ reshape_op_input_tensor = _get_tensor(input_tensor_node)`
`659`	`661`	`assert reshape_op_input_tensor.ndim > reshape_op_output_tensor.ndim, (`
`660`	`662`	`"reshape must be from 3D+ to 2D"`
`661`	`663`	`)`