Update on "[ET-VK] Add pass to remove local_scalar_dense"

## Context

Scalar tensors (i.e. tensors with only 1 element) are often passed in to functions as scalars via

```
scalar_tensor[0].item()
```

This translates to the following chain in the graph

```
index_select = index_select(scalar_tensor, ...)
scalar = local_scalar_dense(index_select)
```

This diff introduces a pass to remove the `local_scalar_dense` "chain" in favor of passing in the input tensor directly.

Note that this replacement only occurs if the original tensor is a scalar tensor.

In the Vulkan backend, these scalar tensors will be represented as symbolic integers instead of actual tensors, which is why this replacement is valid. However, it may not a valid replacement for other backends.

Differential Revision: [D63913432](https://our.internmc.facebook.com/intern/diff/D63913432/)

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/partitioner/supported_ops.py

@@ -48,15 +48,15 @@ def __contains__(self, op):
 ]
 
 SUPPORTS_DYNAMIC_SHAPE = [
-    # Binary broadcasting operators
+    # Binary broadcasting
     exir_ops.edge.aten.add.Tensor,
     exir_ops.edge.aten.sub.Tensor,
     exir_ops.edge.aten.minimum.default,
     exir_ops.edge.aten.mul.Tensor,
     exir_ops.edge.aten.div.Tensor,
     exir_ops.edge.aten.div.Tensor_mode,
     exir_ops.edge.aten.pow.Tensor_Tensor,
-    # Unary elementwise operators
+    # Unary elementwise
     exir_ops.edge.aten.abs.default,
     exir_ops.edge.aten.clamp.default,
     exir_ops.edge.aten.cos.default,
@@ -70,46 +70,46 @@ def __contains__(self, op):
     exir_ops.edge.aten.sin.default,
     exir_ops.edge.aten.sqrt.default,
     exir_ops.edge.aten.tanh.default,
-    # Matrix Multiplication Operators
+    # Matrix Multiplication
     exir_ops.edge.aten.bmm.default,
     exir_ops.edge.aten.mm.default,
     exir_ops.edge.aten.addmm.default,
     exir_ops.edge.aten.linear.default,
-    # Reduction operators
+    # Reduction
     exir_ops.edge.aten._log_softmax.default,
     exir_ops.edge.aten._softmax.default,
-    # 2D Pooling ops
+    # 2D Pooling
     exir_ops.edge.aten.avg_pool2d.default,
     exir_ops.edge.aten.max_pool2d_with_indices.default,
-    # Convolution ops
+    # Convolution
     exir_ops.edge.aten.convolution.default,
     exir_ops.edge.et_vk.conv_with_clamp.default,
 ]
 
 NO_DYNAMIC_SHAPE = [
-    # Reduction operators
+    # Reduction
     exir_ops.edge.aten.mean.dim,
     exir_ops.edge.aten.sum.dim_IntList,
-    # Normalization operators
+    # Normalization
     exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
     exir_ops.edge.aten.native_layer_norm.default,
-    # Shape Manipulation operators
+    # Shape Manipulation
     exir_ops.edge.aten.squeeze_copy.dims,
     exir_ops.edge.aten.unsqueeze_copy.default,
     exir_ops.edge.aten.view_copy.default,
     exir_ops.edge.aten.permute_copy.default,
     exir_ops.edge.aten.t_copy.default,
-    # Indexing and lookup operators
+    # Indexing and lookup
     exir_ops.edge.aten.embedding.default,
     exir_ops.edge.aten.index_select.default,
     exir_ops.edge.aten.select_copy.int,
     exir_ops.edge.aten.slice_copy.Tensor,
-    # Tensor combination operators
+    # Tensor combination
     exir_ops.edge.aten.cat.default,
     exir_ops.edge.aten.split_with_sizes_copy.default,
     exir_ops.edge.aten.split.Tensor,
     exir_ops.edge.aten.repeat.default,
-    # Tensor creation operators
+    # Tensor creation
     exir_ops.edge.aten.arange.start_step,
     exir_ops.edge.aten.clone.default,
     exir_ops.edge.aten.constant_pad_nd.default,