Add complex dtype support to edge dialect for specific ops (#20828)

The edge dialect dtype constraint tables did not include complex dtypes
(ComplexFloat, ComplexDouble), causing the verifier to incorrectly
reject models that use complex tensors. The runtime already supports
these types (scalar_type.h indices 9-10).

Registers ComplexFloat and ComplexDouble as recognized dtype names and
adds them to the following shape-manipulation ops that are
dtype-agnostic: unsqueeze_copy, squeeze_copy, permute_copy,
slice_copy, expand_copy.

This was forcing downstream exporters to monkey-patch
_check_tensor_args_matching_op_allowed_dtype to suppress
SpecViolationError for complex dtype violations.

Author: psiddh

exir/dialects/edge/dtype/supported.py

@@ -27,6 +27,8 @@
     torch.float: "Float",
     torch.double: "Double",
     torch.uint16: "UInt16",
+    torch.complex64: "ComplexFloat",
+    torch.complex128: "ComplexDouble",
 }
 
 regular_tensor_str_to_dtypes = {

exir/dialects/edge/edge.yaml

@@ -3777,7 +3777,7 @@
   namespace: edge
   inherits: aten::expand_copy
   type_alias:
-    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16]
+    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16, ComplexFloat, ComplexDouble]
   type_constraint:
   - self: T0
     __ret_0: T0
@@ -6523,7 +6523,7 @@
   namespace: edge
   inherits: aten::permute_copy
   type_alias:
-    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16]
+    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16, ComplexFloat, ComplexDouble]
   type_constraint:
   - self: T0
     __ret_0: T0
@@ -7185,7 +7185,7 @@
   namespace: edge
   inherits: aten::slice_copy.Tensor
   type_alias:
-    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16]
+    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16, ComplexFloat, ComplexDouble]
   type_constraint:
   - self: T0
     __ret_0: T0
@@ -7272,7 +7272,7 @@
   namespace: edge
   inherits: aten::squeeze_copy.dim
   type_alias:
-    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16]
+    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16, ComplexFloat, ComplexDouble]
   type_constraint:
   - self: T0
     __ret_0: T0
@@ -7281,7 +7281,7 @@
   namespace: edge
   inherits: aten::squeeze_copy.dims
   type_alias:
-    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16]
+    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16, ComplexFloat, ComplexDouble]
   type_constraint:
   - self: T0
     __ret_0: T0
@@ -7821,7 +7821,7 @@
   namespace: edge
   inherits: aten::unsqueeze_copy
   type_alias:
-    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16]
+    T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16, ComplexFloat, ComplexDouble]
   type_constraint:
   - self: T0
     __ret_0: T0

exir/tests/test_verification.py

@@ -273,3 +273,54 @@ def forward(self, x):
                 .exported_program()
                 .graph_module
             )
+
+    def test_edge_unsqueeze_complex_dtype(self) -> None:
+        class TestModel(torch.nn.Module):
+            def forward(self, x):
+                return x.unsqueeze(0)
+
+        m = TestModel()
+        for dtype in (torch.complex64, torch.complex128):
+            edge_gm = (
+                to_edge(
+                    export(
+                        m,
+                        (torch.randn(4, 8).to(dtype=dtype),),
+                        strict=True,
+                    )
+                )
+                .exported_program()
+                .graph_module
+            )
+            self.assertTrue(
+                any("unsqueeze_copy" in str(n.target) for n in edge_gm.graph.nodes),
+                "Expected unsqueeze to lower to unsqueeze_copy",
+            )
+            verifier = EXIREdgeDialectVerifier()
+            verifier(edge_gm)
+            self.assertTrue(verifier.is_valid(edge_gm))
+
+    def test_edge_shape_ops_complex_dtype(self) -> None:
+        class TestModel(torch.nn.Module):
+            def forward(self, x):
+                x = x.expand(4, 8)
+                x = x.permute(1, 0)
+                x = x[:, :2]
+                x = x.squeeze(0)
+                return x
+
+        m = TestModel()
+        edge_gm = (
+            to_edge(
+                export(
+                    m,
+                    (torch.randn(1, 8).to(dtype=torch.complex64),),
+                    strict=True,
+                )
+            )
+            .exported_program()
+            .graph_module
+        )
+        verifier = EXIREdgeDialectVerifier()
+        verifier(edge_gm)
+        self.assertTrue(verifier.is_valid(edge_gm))