Optimise and remove magic constants

Signed-off-by: neNasko1 <nasko119@gmail.com>
Quantco · Jan 7, 2025 · 3e22b77 · 3e22b77
1 parent 6bd7b43
commit 3e22b77
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/ndonnx/_core/_numericimpl.py b/ndonnx/_core/_numericimpl.py
@@ -98,8 +98,13 @@ def bitwise_or(self, x, y):
     @validate_core
     def bitwise_right_shift(self, x, y):
         # Since we need to perform arithmetic right-shift we have to be a bit more careful
+        if isinstance(x.dtype, (dtypes.Unsigned, dtypes.NullableUnsigned)):
+            return binary_op(
+                x, y, lambda a, b: opx.bit_shift(a, b, direction="RIGHT"), dtypes.uint64
+            )
+        MAX_POW = 63
         return ndx.where(
-            y >= 63,
+            y >= MAX_POW,
             ndx.where(x >= 0, 0, -1),
             ndx.floor_divide(x, ndx.pow(ndx.asarray(2, ndx.int64), y)),
         ).astype(x.dtype)

diff --git a/tests/test_core.py b/tests/test_core.py
@@ -1,4 +1,4 @@
-# Copyright (c) QuantCo 2023-2024
+# Copyright (c) QuantCo 2023-2025
 # SPDX-License-Identifier: BSD-3-Clause
 
 from __future__ import annotations