Fix weekly failures

[neNasko1]

This PR fixes the weekly failures of the array-api(#92) by:

• disabling argmin and argmax on the test suite and improving our tests as onnxruntime does not implement them for UInt64
• reimplements right shift as the array-api defines it as an arithmetic right shift
• fixes all and any for non-boolean inputs
• raises errors when possible in expand_dims and getitem_null

[cbourjau]

Please note that where executes both branches so you may still run into a division by zero issue.

[adityagoel4512]

Which division by zero are you referring to?

[cbourjau]

Line 109 due to an overflow.

[adityagoel4512]

Fair point, I suspected the same but second-guessed it when the test looked like it passed! The explanation for why it silently succeed looks to be that ndonnx (or really onnxruntime) and NumPy overflowed quite differently :/.

(Pdb) np.__version__
'2.1.3'
(Pdb) ndx.__version__
'0.0.post75+g3e22b77'
(Pdb) np.pow(np.asarray(2, np.int64), 70)
np.int64(0)
(Pdb) ndx.pow(ndx.asarray(2, ndx.int64), 70)
Array(9223372036854775807, dtype=Int64)

But yes, let's not depend on this @neNasko1. I'm pretty sure you can just extract the sign bit and manually roll this with the logical right shift. Additionally, this should simply not work for floating point arrays. We'll improve this with typed arrays but for now you'll need to guard that.

I think the following should help avoid division.

-        MAX_POW = 63
-        return ndx.where(
-            y >= MAX_POW,
-            ndx.where(x >= 0, 0, -1),
-            ndx.floor_divide(x, ndx.pow(ndx.asarray(2, ndx.int64), y)),
-        ).astype(x.dtype)
+        elif isinstance(x.dtype, (dtypes.Integral, dtypes.NullableIntegral)):
+            mask = ndx.where(
+                x >= 0,
+                ndx.asarray(0, dtype=x.dtype),
+                (((ndx.asarray(1, dtype=ndx.uint64) << y) - 1) << (ndx.iinfo(x.dtype).bits-y)).astype(x.dtype)
+            )
+            logic_shift = binary_op(
+                x, y, lambda a, b: opx.bit_shift(a, b, direction="RIGHT"), dtypes.uint64
+            ) | mask
+            MAX_POW = 63
+            return ndx.where(
+                y >= MAX_POW,
+                ndx.where(x >= 0, 0, -1),
+                logic_shift,
+            ).astype(x.dtype)
+        else:
+            return NotImplemented

[neNasko1]

I rewrote the solution a bit, can you take a look?

[adityagoel4512]

Fair point, I suspected the same but second-guessed it when the test looked like it passed! The explanation for why it silently succeed looks to be that ndonnx (or really onnxruntime) and NumPy overflowed quite differently :/.

(Pdb) np.__version__
'2.1.3'
(Pdb) ndx.__version__
'0.0.post75+g3e22b77'
(Pdb) np.pow(np.asarray(2, np.int64), 70)
np.int64(0)
(Pdb) ndx.pow(ndx.asarray(2, ndx.int64), 70)
Array(9223372036854775807, dtype=Int64)

But yes, let's not depend on this @neNasko1. I'm pretty sure you can just extract the sign bit and manually roll this with the logical right shift. Additionally, this should simply not work for floating point arrays. We'll improve this with typed arrays but for now you'll need to guard that.

I think the following should help avoid division.

-        MAX_POW = 63
-        return ndx.where(
-            y >= MAX_POW,
-            ndx.where(x >= 0, 0, -1),
-            ndx.floor_divide(x, ndx.pow(ndx.asarray(2, ndx.int64), y)),
-        ).astype(x.dtype)
+        elif isinstance(x.dtype, (dtypes.Integral, dtypes.NullableIntegral)):
+            mask = ndx.where(
+                x >= 0,
+                ndx.asarray(0, dtype=x.dtype),
+                (((ndx.asarray(1, dtype=ndx.uint64) << y) - 1) << (ndx.iinfo(x.dtype).bits-y)).astype(x.dtype)
+            )
+            logic_shift = binary_op(
+                x, y, lambda a, b: opx.bit_shift(a, b, direction="RIGHT"), dtypes.uint64
+            ) | mask
+            MAX_POW = 63
+            return ndx.where(
+                y >= MAX_POW,
+                ndx.where(x >= 0, 0, -1),
+                logic_shift,
+            ).astype(x.dtype)
+        else:
+            return NotImplemented

[adityagoel4512]

Could you please add a short test to ensure we IndexError in the correct situations with boolean indexing?

[neNasko1]

Done!