Merge pull request #247 from magnatelee/ufunc-in-metamethods

Use ufuncs in special methods
nv-legate · Mar 31, 2022 · faa7b25 · faa7b25
2 parents 1b70202 + c77cf6b
commit faa7b25
Show file tree

Hide file tree

Showing 2 changed files with 155 additions and 292 deletions.
diff --git a/cunumeric/_ufunc/ufunc.py b/cunumeric/_ufunc/ufunc.py
@@ -132,7 +132,7 @@ def _maybe_cast_input(self, arr, to_dtype, casting):
         if arr.dtype == to_dtype:
             return arr
 
-        if not np_can_cast(arr.dtype, to_dtype):
+        if not np_can_cast(arr.dtype, to_dtype, casting=casting):
             raise TypeError(
                 f"Cannot cast ufunc '{self._name}' input from "
                 f"{arr.dtype} to {to_dtype} with casting rule '{casting}'"
@@ -244,7 +244,7 @@ def __call__(
             out = result
         else:
             if out.dtype != res_dtype:
-                if not np_can_cast(res_dtype, out.dtype):
+                if not np_can_cast(res_dtype, out.dtype, casting=casting):
                     raise TypeError(
                         f"Cannot cast ufunc '{self._name}' output from "
                         f"{res_dtype} to {out.dtype} with casting rule "
@@ -393,7 +393,7 @@ def __call__(
             out = result
         else:
             if out.dtype != res_dtype:
-                if not np_can_cast(res_dtype, out.dtype):
+                if not np_can_cast(res_dtype, out.dtype, casting=casting):
                     raise TypeError(
                         f"Cannot cast ufunc '{self._name}' output from "
                         f"{res_dtype} to {out.dtype} with casting rule "