Fix pow cast (#39)

Quantco · Aug 1, 2024 · 6d01494 · 6d01494
1 parent a801af1
commit 6d01494
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/ndonnx/_core/_impl.py b/ndonnx/_core/_impl.py
@@ -289,7 +289,12 @@ def positive(self, x):
     def pow(self, x, y):
         x, y = ndx.asarray(x), ndx.asarray(y)
         dtype = ndx.result_type(x, y)
-        if isinstance(dtype, (dtypes.Unsigned, dtypes.NullableUnsigned)):
+        if isinstance(dtype, (dtypes.Unsigned, dtypes.NullableUnsigned)) or dtype in (
+            dtypes.int8,
+            dtypes.nint8,
+            dtypes.int16,
+            dtypes.nint16,
+        ):
             return _binary_op(x, y, opx.pow, dtypes.int64)
         else:
             return _binary_op(x, y, opx.pow)