Closes Bears-R-Us#3954: skip_by_rank to handle set containment

pytest.ini

@@ -74,5 +74,6 @@ env =
 markers =
     skip_if_max_rank_less_than
     skip_if_max_rank_greater_than
+    skip_if_rank_not_compiled
     skip_if_nl_greater_than
     skip_if_nl_less_than

tests/array_api/array_creation.py

@@ -15,21 +15,21 @@
 
 
 class TestArrayCreation:
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     @pytest.mark.parametrize("shape", SHAPES)
     @pytest.mark.parametrize("dtype", ak.ScalarDTypes)
     def test_zeros(self, shape, dtype):
         a = xp.zeros(shape, dtype=dtype)
         assert_equivalent(a._array, np.zeros(shape, dtype=dtype))
 
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     @pytest.mark.parametrize("shape", SHAPES)
     @pytest.mark.parametrize("dtype", ak.ScalarDTypes)
     def test_ones(self, shape, dtype):
         a = xp.ones(shape, dtype=dtype)
         assert_equivalent(a._array, np.ones(shape, dtype=dtype))
 
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     def test_from_numpy(self):
         # TODO: support 0D (scalar) arrays
         # (need changes to the create0D command from #2967)
@@ -41,7 +41,7 @@ def test_from_numpy(self):
             assert b.shape == a.shape
             assert b.tolist() == a.tolist()
 
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     @pytest.mark.parametrize("data_type", [ak.int64, ak.float64, ak.bool_])
     @pytest.mark.parametrize("prob_size", pytest.prob_size)
     def test_triu(self, data_type, prob_size):
@@ -60,7 +60,7 @@ def test_triu(self, data_type, prob_size):
                 ak_triu = array_triu(pda, k=diag)._array
                 assert_almost_equivalent(ak_triu, np_triu)
 
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     @pytest.mark.parametrize("data_type", [ak.int64, ak.float64, ak.bool_])
     @pytest.mark.parametrize("prob_size", pytest.prob_size)
     def test_tril(self, data_type, prob_size):
@@ -79,7 +79,7 @@ def test_tril(self, data_type, prob_size):
                 ak_tril = array_tril(pda, k=diag)._array
                 assert_almost_equivalent(np_tril, ak_tril)
 
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     @pytest.mark.parametrize("data_type", [ak.int64, ak.float64, ak.bool_])
     @pytest.mark.parametrize("prob_size", pytest.prob_size)
     def test_eye(self, data_type, prob_size):

tests/array_api/array_manipulation.py

@@ -18,7 +18,7 @@ def randArr(shape):
 
 class TestManipulation:
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_broadcast(self):
         a = xp.ones((1, 6, 1))
         b = xp.ones((5, 1, 10))
@@ -37,7 +37,7 @@ def test_broadcast(self):
         assert (abcd[2] == 1).all()
         assert (abcd[3] == 1).all()
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_concat(self):
         a = randArr((5, 3, 10))
         b = randArr((5, 3, 2))
@@ -70,7 +70,7 @@ def test_concat(self):
         assert hijConcat.shape == (18,)
         assert hijConcat.tolist() == hijNP.tolist()
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_expand_dims(self):
         a = randArr((5, 3))
         alist = a.tolist()
@@ -105,7 +105,7 @@ def test_expand_dims(self):
         with pytest.raises(IndexError):
             xp.expand_dims(a, axis=-4)
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([3])
     def test_flip(self):
         # 1D case
         a = xp.arange(10)
@@ -145,7 +145,7 @@ def test_flip(self):
         with pytest.raises(IndexError):
             xp.flip(r, axis=-4)
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([3])
     def test_permute_dims(self):
         r = randArr((7, 8, 9))
 
@@ -172,7 +172,7 @@ def test_permute_dims(self):
         with pytest.raises(IndexError):
             xp.permute_dims(r, (0, 1, -4))
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_reshape(self):
         r = randArr((2, 6, 12))
         nr = np.asarray(r.tolist())
@@ -201,7 +201,7 @@ def test_reshape(self):
             # more than one dimension can't be inferred
             xp.reshape(r, (2, -1, -1))
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([3])
     def test_roll(self):
         # 1D case
         a = xp.arange(10)
@@ -244,7 +244,7 @@ def test_roll(self):
         with pytest.raises(IndexError):
             xp.roll(r, 3, axis=-4)
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_squeeze(self):
         r1 = randArr((1, 2, 3))
         r2 = randArr((2, 1, 3))
@@ -277,7 +277,7 @@ def test_squeeze(self):
         with pytest.raises(ValueError):
             xp.squeeze(r4, axis=1)
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_stack_unstack(self):
         a = randArr((5, 4))
         b = randArr((5, 4))
@@ -303,7 +303,7 @@ def test_stack_unstack(self):
         assert bp.tolist() == b.tolist()
         assert cp.tolist() == c.tolist()
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_tile(self):
         a = randArr((2, 3))
 
@@ -313,7 +313,7 @@ def test_tile(self):
             assert at.shape == npat.shape
             assert at.tolist() == npat.tolist()
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     def test_repeat(self):
         a = randArr((5, 10))
         r = randArr((50,))

tests/array_api/array_object.py

@@ -13,7 +13,7 @@ def test_chunk_info(self):
         assert chunks == [[0]]
 
     @pytest.mark.skipif(pytest.nl > 1, reason="Multi-local will produce different chunk_info")
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     def test_chunk_info_2dim(self):
 
         a = xp.zeros((2, 2))

tests/array_api/binary_ops.py

@@ -15,7 +15,7 @@
 
 
 class TestArrayCreation:
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     @pytest.mark.parametrize("op", ["+", "-", "*", "/"])
     @pytest.mark.parametrize("dtype", SCALAR_TYPES)
     def test_binops(self, op, dtype):

tests/array_api/elementwise_functions.py

@@ -11,7 +11,7 @@
 
 
 class TestElemenwiseFunctions:
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     def test_logical_not(self):
         a = xp.asarray(ak.array([True, False, True, False]))
         not_a = xp.asarray(ak.array([False, True, False, True]))

tests/array_api/indexing.py

@@ -17,7 +17,7 @@ def randArr(shape):
 
 
 class TestIndexing:
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_rank_changing_assignment(self):
         a = randArr((5, 6, 7))
         b = randArr((5, 6))
@@ -37,7 +37,7 @@ def test_rank_changing_assignment(self):
         a[:, :, :] = e
         assert a.tolist() == e.tolist()
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([3])
     def test_nd_assignment(self):
         a = randArr((5, 6, 7))
         bnp = randArr((5, 6, 7)).to_ndarray()
@@ -51,7 +51,7 @@ def test_nd_assignment(self):
         a[:] = 5
         assert (a == 5).all()
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_pdarray_index(self):
         a = randArr((5, 6, 7))
         anp = np.asarray(a.tolist())
@@ -82,7 +82,7 @@ def test_pdarray_index(self):
         xnp = anp[:]
         assert x.tolist() == xnp.tolist()
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_none_index(self):
         a = randArr((10, 10))
         anp = np.asarray(a.tolist())

tests/array_api/linalg.py

@@ -8,7 +8,7 @@
 
 class TestLinalg:
 
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     @pytest.mark.parametrize("data_type1", [ak.int64, ak.float64, ak.bool_])
     @pytest.mark.parametrize("data_type2", [ak.int64, ak.float64, ak.bool_])
     @pytest.mark.parametrize("prob_size", pytest.prob_size)
@@ -27,7 +27,7 @@ def test_matmul(self, data_type1, data_type2, prob_size):
             npProduct = np.matmul(ndaLeft, ndaRight)
             assert_almost_equivalent(akProduct._array, npProduct)
 
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     @pytest.mark.parametrize("data_type", [ak.int64, ak.float64, ak.bool_])
     @pytest.mark.parametrize("prob_size", pytest.prob_size)
     def test_transpose(self, data_type, prob_size):
@@ -43,7 +43,7 @@ def test_transpose(self, data_type, prob_size):
             ppa = xp.matrix_transpose(array)._array
             assert np.allclose(ppa.to_ndarray(), npa)
 
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     @pytest.mark.parametrize("data_type1", [ak.int64, ak.float64, ak.bool_])
     @pytest.mark.parametrize("data_type2", [ak.int64, ak.float64])
     @pytest.mark.parametrize("prob_size", pytest.prob_size)

tests/array_api/searching_functions.py

@@ -10,7 +10,7 @@
 
 
 class TestSearchingFunctions:
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_argmax(self):
         a = xp.asarray(ak.randint(0, 100, (4, 5, 6), dtype=ak.int64, seed=SEED))
         a[3, 2, 1] = 101
@@ -27,7 +27,7 @@ def test_argmax(self):
         assert aArgmax1Keepdims.shape == (4, 1, 6)
         assert aArgmax1Keepdims[3, 0, 1] == 2
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_argmin(self):
         a = xp.asarray(ak.randint(0, 100, (4, 5, 6), dtype=ak.int64, seed=SEED))
         a[3, 2, 1] = -1
@@ -42,7 +42,7 @@ def test_argmin(self):
         assert aArgmin1Keepdims.shape == (4, 1, 6)
         assert aArgmin1Keepdims[3, 0, 1] == 2
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([3])
     def test_nonzero(self):
         a = xp.zeros((40, 15, 16), dtype=ak.int64)
         a[0, 1, 0] = 1
@@ -75,7 +75,7 @@ def test_nonzero_1d(self):
 
         assert nz[0].tolist() == [0, 12, 100, 205, 490]
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([3])
     def test_where(self):
         a = xp.zeros((4, 5, 6), dtype=ak.int64)
         a[1, 2, 3] = 1
@@ -94,7 +94,7 @@ def test_where(self):
         assert d[0, 0, 0] == c[0, 0, 0]
         assert d[3, 3, 3] == c[3, 3, 3]
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     def test_search_sorted(self):
         a = xp.asarray(ak.randint(0, 100, 1000, dtype=ak.float64))
         b = xp.asarray(ak.randint(0, 100, (10, 10), dtype=ak.float64))

tests/array_api/set_functions.py

@@ -10,9 +10,9 @@
 
 def ret_shapes():
     shapes = [1000]
-    if pytest.max_rank > 1:
+    if 2 in ak.client.get_array_ranks():
         shapes.append((20, 50))
-    if pytest.max_rank > 2:
+    if 3 in ak.client.get_array_ranks():
         shapes.append((2, 10, 50))
     return shapes
 

tests/array_api/sorting.py

@@ -14,7 +14,7 @@
 
 class TestArrayCreation:
 
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     def test_argsort(self):
         for shape in SHAPES:
             for dtype in ak.ScalarDTypes:
@@ -46,7 +46,7 @@ def test_argsort(self):
                                 for j in range(shape[1] - 1):
                                     assert a[i, b[i, j]] <= a[i, b[i, j + 1]]
 
-    @pytest.mark.skip_if_max_rank_less_than(2)
+    @pytest.mark.skip_if_rank_not_compiled([2])
     @pytest.mark.parametrize("dtype", SCALAR_TYPES)
     @pytest.mark.parametrize("shape", SHAPES)
     def test_sort(self, dtype, shape):

tests/array_api/stats_functions.py

@@ -12,7 +12,7 @@
 
 class TestStatsFunction:
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_max(self):
         a = xp.asarray(ak.randint(0, 100, (5, 7, 4), dtype=ak.int64, seed=SEED))
         a[3, 6, 2] = 101
@@ -31,7 +31,7 @@ def test_max(self):
         assert aMax02Keepdims.shape == (1, 7, 1)
         assert aMax02Keepdims[0, 6, 0] == 101
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_min(self):
         a = xp.asarray(ak.randint(0, 100, (5, 7, 4), dtype=ak.int64, seed=SEED))
         a[3, 6, 2] = -1
@@ -50,7 +50,7 @@ def test_min(self):
         assert aMin02Keepdims.shape == (1, 7, 1)
         assert aMin02Keepdims[0, 6, 0] == -1
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_mean(self):
         a = xp.ones((10, 5, 5))
         a[0, 0, 0] = 251
@@ -76,7 +76,7 @@ def test_mean(self):
         assert aMean02Keepdims[0, 0, 0] == 2
         assert aMean02Keepdims[2, 0, 0] == 2
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_std(self):
         a = xp.ones((10, 5, 5), dtype=ak.float64)
         a[0, 0, 0] = 26
@@ -96,7 +96,7 @@ def test_std(self):
         assert abs(aStd02Keepdims[0, 0, 0] - math.sqrt(24)) < 1e-10
         assert abs(aStd02Keepdims[2, 0, 0] - math.sqrt(24)) < 1e-10
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_var(self):
         a = xp.ones((10, 5, 5), dtype=ak.float64)
         a[0, 0, 0] = 26
@@ -116,7 +116,7 @@ def test_var(self):
         assert abs(aStd02Keepdims[0, 0, 0] - 24) < 1e-10
         assert abs(aStd02Keepdims[2, 0, 0] - 24) < 1e-10
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_prod(self):
         a = xp.ones((2, 3, 4))
         a = a + a
@@ -135,7 +135,7 @@ def test_prod(self):
         assert aProd02Keepdims.shape == (2, 1, 1)
         assert aProd02Keepdims[0, 0, 0] == 2**12
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([2, 3])
     def test_sum(self):
         a = xp.ones((2, 3, 4))
 
@@ -153,7 +153,7 @@ def test_sum(self):
         assert aSum02Keepdims.shape == (2, 1, 1)
         assert aSum02Keepdims[0, 0, 0] == 12
 
-    @pytest.mark.skip_if_max_rank_less_than(3)
+    @pytest.mark.skip_if_rank_not_compiled([3])
     def test_cumsum(self):
         a = xp.asarray(ak.randint(0, 100, (5, 6, 7), seed=SEED))