Add a pass to convert rank-0 tensor to rank-1 tensor

Differential Revision: D69281867

Author: cccclai

backends/transforms/rank_0_to_rank_1.py

@@ -0,0 +1,17 @@
+import torch
+from executorch.exir.pass_base import ExportPass, PassResult
+
+class Rank0ToRank1Pass(ExportPass):
+    """
+    Replace Rank-0 Tensor to Rank-1 Tensor for all the inputs.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        for node in graph_module.graph.nodes:
+            if node.op == "placeholder":
+                node.meta["val"] = node.meta["val"].reshape(1, 1)
+        graph_module.recompile()
+        return PassResult(graph_module, True)

backends/transforms/targets.bzl

@@ -187,6 +187,20 @@ def define_common_targets():
         ],
     )
 
+    runtime.python_library(
+        name = "rank_0_to_rank_1",
+        srcs = [
+            "rank_0_to_rank_1.py",
+        ],
+        visibility = [
+            "//executorch/backends/...",
+        ],
+        deps = [
+            "//caffe2:torch",
+            "//executorch/exir:pass_base",
+        ],
+    )
+
     runtime.python_test(
         name = "test_duplicate_dynamic_quant_chain",
         srcs = [