Arm backend: Squeeze on the dimension we have selected (#8335)

SaoirseARM · web-flow · commit 9746ce709654 · 2025-02-10T15:54:55.000+01:00
Squeeze on the dimension we have selected

- If we squeeze using the input rank and batch_size==1 this dimension will also be removed.
diff --git a/backends/arm/_passes/decompose_select.py b/backends/arm/_passes/decompose_select.py
@@ -37,14 +37,13 @@ def call(self, graph_module: torch.fx.GraphModule):
             rank = len(input_node.meta["val"].size())
             dim = dim % rank if dim < 0 else dim
             index = index % rank if index < 0 else index
-            dim_list = list(range(rank))
 
             with graph_module.graph.inserting_before(node):
                 slice_node = create_node(
                     graph_module.graph, slice_op, (input_node, dim, index, index + 1)
                 )
                 squeeze_node = create_node(
-                    graph_module.graph, squeeze_op, (slice_node, dim_list)
+                    graph_module.graph, squeeze_op, (slice_node, [dim])
                 )
 
             node.replace_all_uses_with(squeeze_node)

Original file line number	Diff line number	Diff line change
`@@ -37,14 +37,13 @@ def call(self, graph_module: torch.fx.GraphModule):`
`37`	`37`	`rank = len(input_node.meta["val"].size())`
`38`	`38`	`dim = dim % rank if dim < 0 else dim`
`39`	`39`	`index = index % rank if index < 0 else index`
`40`		`- dim_list = list(range(rank))`
`41`	`40`
`42`	`41`	`with graph_module.graph.inserting_before(node):`
`43`	`42`	`slice_node = create_node(`
`44`	`43`	`graph_module.graph, slice_op, (input_node, dim, index, index + 1)`
`45`	`44`	`)`
`46`	`45`	`squeeze_node = create_node(`
`47`		`- graph_module.graph, squeeze_op, (slice_node, dim_list)`
	`46`	`+ graph_module.graph, squeeze_op, (slice_node, [dim])`
`48`	`47`	`)`
`49`	`48`
`50`	`49`	`node.replace_all_uses_with(squeeze_node)`