add complex

Author: xadupre

onnx_array_api/_helpers.py

@@ -40,6 +40,10 @@ def np_dtype_to_tensor_dtype(dtype: Any):
             dt = TensorProto.INT64
         elif dtype is float:
             dt = TensorProto.DOUBLE
+        elif dtype == np.complex64:
+            dt = TensorProto.COMPLEX64
+        elif dtype == np.complex128:
+            dt = TensorProto.COMPLEX128
         else:
             raise KeyError(f"Unable to guess type for dtype={dtype}.")  # noqa: B904
     return dt

onnx_array_api/annotations.py

@@ -64,6 +64,8 @@ def wrapper(self, *args: List[Any], **kwargs: Dict[str, Any]) -> Any:
     np.uint64: TensorProto.UINT64,
     np.bool_: TensorProto.BOOL,
     np.str_: TensorProto.STRING,
+    np.complex64: TensorProto.COMPLEX64,
+    np.complex128: TensorProto.COMPLEX128,
 }
 
 

onnx_array_api/array_api/__init__.py

@@ -47,6 +47,8 @@ def _finfo(dtype):
             continue
         if isinstance(v, (np.float32, np.float64, np.float16)):
             d[k] = float(v)
+        elif isinstance(v, (np.complex128, np.complex64)):
+            d[k] = complex(v)
         else:
             d[k] = v
     d["dtype"] = DType(np_dtype_to_tensor_dtype(dt))
@@ -124,6 +126,8 @@ def _finalize_array_api(module, function_names, TEagerTensor):
     module.float16 = DType(TensorProto.FLOAT16)
     module.float32 = DType(TensorProto.FLOAT)
     module.float64 = DType(TensorProto.DOUBLE)
+    module.complex64 = DType(TensorProto.COMPLEX64)
+    module.complex128 = DType(TensorProto.COMPLEX128)
     module.int8 = DType(TensorProto.INT8)
     module.int16 = DType(TensorProto.INT16)
     module.int32 = DType(TensorProto.INT32)

onnx_array_api/npx/npx_var.py

@@ -1171,13 +1171,17 @@ def __init__(self, cst: Any):
             Var.__init__(self, np.array(cst, dtype=np.int64), op="Identity")
         elif isinstance(cst, float):
             Var.__init__(self, np.array(cst, dtype=np.float64), op="Identity")
+        elif isinstance(cst, complex):
+            Var.__init__(self, np.array(cst, dtype=np.complex128), op="Identity")
         elif isinstance(cst, list):
             if all(isinstance(t, bool) for t in cst):
                 Var.__init__(self, np.array(cst, dtype=np.bool_), op="Identity")
             elif all(isinstance(t, (int, bool)) for t in cst):
                 Var.__init__(self, np.array(cst, dtype=np.int64), op="Identity")
             elif all(isinstance(t, (float, int, bool)) for t in cst):
                 Var.__init__(self, np.array(cst, dtype=np.float64), op="Identity")
+            elif all(isinstance(t, (float, int, bool, complex)) for t in cst):
+                Var.__init__(self, np.array(cst, dtype=np.complex128), op="Identity")
             else:
                 raise ValueError(
                     f"Unable to convert cst (type={type(cst)}), value={cst}."

onnx_array_api/reference/evaluator_yield.py

@@ -485,6 +485,12 @@ def generate_input(info: ValueInfoProto) -> np.ndarray:
         return (value.astype(np.float16) / p).astype(np.float16).reshape(new_shape)
     if elem_type == TensorProto.DOUBLE:
         return (value.astype(np.float64) / p).astype(np.float64).reshape(new_shape)
+    if elem_type == TensorProto.COMPLEX64:
+        return (value.astype(np.complex64) / p).astype(np.complex64).reshape(new_shape)
+    if elem_type == TensorProto.COMPLEX128:
+        return (
+            (value.astype(np.complex128) / p).astype(np.complex128).reshape(new_shape)
+        )
     raise RuntimeError(f"Unexpected element_type {elem_type} for info={info}")
 
 

onnx_array_api/reference/ops/op_constant_of_shape.py

@@ -19,6 +19,8 @@ def _process(value):
             cst = np.int64(cst)
         elif isinstance(cst, float):
             cst = np.float64(cst)
+        elif isinstance(cst, complex):
+            cst = np.complex128(cst)
         elif cst is None:
             cst = np.float32(0)
         if not isinstance(
@@ -27,6 +29,8 @@ def _process(value):
                 np.float16,
                 np.float32,
                 np.float64,
+                np.complex64,
+                np.complex128,
                 np.int64,
                 np.int32,
                 np.int16,