From 7760e8eac3286aa93771fa02cb3ab5648695e6f5 Mon Sep 17 00:00:00 2001 From: Aaron Meurer Date: Mon, 5 Jun 2023 14:33:25 -0600 Subject: [PATCH] Add complex dtypes to numpy.array_api result_type --- numpy/array_api/_dtypes.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/numpy/array_api/_dtypes.py b/numpy/array_api/_dtypes.py index 476d619fee63..06acab8e80df 100644 --- a/numpy/array_api/_dtypes.py +++ b/numpy/array_api/_dtypes.py @@ -12,6 +12,8 @@ uint64 = np.dtype("uint64") float32 = np.dtype("float32") float64 = np.dtype("float64") +complex64 = np.dtype("complex64") +complex128 = np.dtype("complex128") # Note: This name is changed bool = np.dtype("bool") @@ -26,10 +28,13 @@ uint64, float32, float64, + complex64, + complex128, bool, ) _boolean_dtypes = (bool,) -_floating_dtypes = (float32, float64) +_real_floating_dtypes = (float32, float64) +_floating_dtypes = (float32, float64, complex64, complex128) _integer_dtypes = (int8, int16, int32, int64, uint8, uint16, uint32, uint64) _integer_or_boolean_dtypes = ( bool, @@ -45,6 +50,8 @@ _numeric_dtypes = ( float32, float64, + complex64, + complex128, int8, int16, int32, @@ -61,6 +68,7 @@ "integer": _integer_dtypes, "integer or boolean": _integer_or_boolean_dtypes, "boolean": _boolean_dtypes, + "real floating-point": _floating_dtypes, "floating-point": _floating_dtypes, } @@ -133,6 +141,18 @@ (float32, float64): float64, (float64, float32): float64, (float64, float64): float64, + (complex64, complex64): complex64, + (complex64, complex128): complex128, + (complex128, complex64): complex128, + (complex128, complex64): complex128, + (float32, complex64): complex64, + (float32, complex128): complex128, + (float64, complex64): complex128, + (float64, complex128): complex128, + (complex64, float32): complex64, + (complex64, float64): complex128, + (complex128, float32): complex128, + (complex128, float64): complex128, (bool, bool): bool, }