Supports complex data type for the add op. (#2061)

Supports complex data type for the `add` op.

Author: junpeiz

coremltools/converters/mil/frontend/torch/ops.py

@@ -869,10 +869,17 @@ def add(context, node):
     # rdar://60175736
     if len(add_inputs) > 2 and add_inputs[2].val != 1:
         raise ValueError("ADD does not support scale factor param")
-    x, y = promote_input_dtypes(add_inputs[:2])
+    x, y = add_inputs[:2]
     if types.is_bool(x.dtype) and types.is_bool(y.dtype):
         add_node = mb.logical_or(x=x, y=y, name=node.name)
+    elif types.is_complex(x.dtype) or types.is_complex(y.dtype):
+        x_real = mb.complex_real(data=x) if types.is_complex(x.dtype) else x
+        x_imag = mb.complex_imag(data=x) if types.is_complex(x.dtype) else 0.0
+        y_real = mb.complex_real(data=y) if types.is_complex(y.dtype) else y
+        y_imag = mb.complex_imag(data=y) if types.is_complex(y.dtype) else 0.0
+        add_node = mb.complex(real_data=mb.add(x=x_real, y=y_real), imag_data=mb.add(x=x_imag, y=y_imag), name=node.name)
     else:
+        x, y = promote_input_dtypes([x, y])
         add_node = mb.add(x=x, y=y, name=node.name)
     context.add(add_node)
 

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -3752,6 +3752,31 @@ def forward(self, x, y):
             input_as_shape=False,
         )
 
+    @pytest.mark.parametrize(
+        "compute_unit, backend, x_complex, y_complex",
+        itertools.product(
+            compute_units,
+            backends,
+            (True, False),
+            (True, False),
+        ),
+    )
+    def test_add_complex(self, compute_unit, backend, x_complex, y_complex):
+        class TestAddComplexModel(nn.Module):
+            def forward(self, x, y):
+                if x_complex:
+                    x = torch.complex(x, x)
+                if y_complex:
+                    y = torch.complex(y, y)
+                return torch.add(x, y).abs()
+
+        self.run_compare_torch(
+            [(2, 3), (2, 3)],
+            TestAddComplexModel(),
+            compute_unit=compute_unit,
+            backend=backend,
+        )
+
 
 class TestFull(TorchBaseTest):
     @pytest.mark.parametrize(