Cast back to input dtype in where

Quantco · Jul 22, 2024 · 8b25df2 · 8b25df2
1 parent 7650f9c
commit 8b25df2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ndonnx/_core/_impl.py b/ndonnx/_core/_impl.py
@@ -520,7 +520,7 @@ def where_dtype_agnostic(a: ndx.Array, b: ndx.Array) -> ndx.Array:
                         a.astype(dtypes.int32)._core(),
                         b.astype(dtypes.int32)._core(),
                     )
-                )
+                ).astype(a.dtype)
             elif isinstance(a.dtype, dtypes.CoreType):
                 return _from_corearray(
                     opx.where(