More tests for all operations

python/mops/src/mops/opsa.py

@@ -11,21 +11,6 @@ def outer_product_scatter_add(A, B, P, n_O):
     A = np.ascontiguousarray(A)
     B = np.ascontiguousarray(B)
     P = np.ascontiguousarray(P)
-
-    # TODO: Include these checks in check_opsa
-    if A.dtype != B.dtype:
-        raise TypeError("A and B must have the same dtype")
-    if len(A.shape) != 2 or len(B.shape) != 2:
-        raise TypeError("A and B must be 2-dimensional arrays")
-    if not np.can_cast(P, np.int32, "same_kind"):
-        raise TypeError("P must be an array of integers")
-    if len(P.shape) != 1:
-        raise TypeError("P must be 1-dimensional arrays")
-    if A.shape[0] != B.shape[0] or A.shape[0] != P.shape[0]:
-        raise TypeError(
-            "A, B and P must have the same number of elements on the " "first dimension"
-        )
-
     P = P.astype(np.int32)
 
     output = np.empty((n_O, A.shape[1], B.shape[1]), dtype=A.dtype)  # TODO: 3D arrays

python/mops/tests/opsa.py

@@ -2,35 +2,77 @@
 import pytest
 from mops.reference_implementations import outer_product_scatter_add as ref_opsa
 
+import mops
 from mops import outer_product_scatter_add as opsa
 
 np.random.seed(0xDEADBEEF)
 
 
-def test_opsa():
+@pytest.fixture
+def valid_arguments():
     A = np.random.rand(100, 20)
     B = np.random.rand(100, 5)
+    P = np.sort(np.random.randint(10, size=(100)))
+    n_O = 20
+    return A, B, P, n_O
 
-    indices = np.sort(np.random.randint(10, size=(100,)))
 
-    reference = ref_opsa(A, B, indices, np.max(indices) + 1)
-    actual = opsa(A, B, indices, np.max(indices) + 1)
-    assert np.allclose(reference, actual)
+def test_opsa(valid_arguments):
+    A, B, P, n_O = valid_arguments
 
+    reference = ref_opsa(A, B, P, n_O)
+    actual = opsa(A, B, P, n_O)
+    assert np.allclose(reference, actual)
 
-def test_opsa_no_neighbors():
-    A = np.random.rand(100, 20)
-    B = np.random.rand(100, 5)
 
-    indices = np.sort(np.random.randint(10, size=(100,)))
+def test_opsa_no_neighbors(valid_arguments):
+    A, B, P, n_O = valid_arguments
     # substitute all 1s by 2s so as to test the no-neighbor case
-    indices[indices == 1] = 2
+    P[P == 1] = 2
 
-    reference = ref_opsa(A, B, indices, np.max(indices) + 1)
-    actual = opsa(A, B, indices, np.max(indices) + 1)
+    reference = ref_opsa(A, B, P, n_O)
+    actual = opsa(A, B, P, n_O)
     assert np.allclose(reference, actual)
 
 
-def test_opsa_wrong_type():
-    with pytest.raises(ValueError):
-        opsa(np.array([1]), 2, 3, 4)
+def test_opsa_wrong_type(valid_arguments):
+    A, B, P, n_O = valid_arguments
+    A = A.astype(np.int32)
+
+    with pytest.raises(TypeError, match="Wrong dtype for A in opsa: got int32"):
+        opsa(A, B, P, n_O)
+
+
+def test_opsa_wrong_number_of_dimensions(valid_arguments):
+    A, B, P, n_O = valid_arguments
+    A = A[..., np.newaxis]
+
+    with pytest.raises(
+        ValueError, match="A must be a 2D array in opsa, got a 3D array"
+    ):
+        opsa(A, B, P, n_O)
+
+
+def test_opsa_size_mismatch(valid_arguments):
+    A, B, P, n_O = valid_arguments
+    P = P[:5]
+
+    with pytest.raises(
+        mops.status.MopsError,
+        match="Dimension mismatch: the sizes of A along "
+        "dimension 0 and P along dimension 0 must match in opsa",
+    ):
+        opsa(A, B, P, n_O)
+
+
+def test_opsa_out_of_bounds(valid_arguments):
+    A, B, P, n_O = valid_arguments
+    P[0] = A.shape[1]
+
+    with pytest.raises(
+        mops.status.MopsError,
+        match="Index array P in operation opsa contains elements up to 20; "
+        "this would cause out-of-bounds accesses. With the provided "
+        "parameters, it can only contain elements up to 19",
+    ):
+        opsa(A, B, P, n_O)

python/mops/tests/opsax.py

@@ -4,24 +4,79 @@
     outer_product_scatter_add_with_weights as ref_opsax,
 )
 
+import mops
 from mops import outer_product_scatter_add_with_weights as opsax
 
 np.random.seed(0xDEADBEEF)
 
 
-def test_opsax():
+@pytest.fixture
+def valid_arguments():
     A = np.random.rand(100, 10)
     R = np.random.rand(100, 5)
     X = np.random.rand(20, 5)
-
     I = np.sort(np.random.randint(20, size=(100,)))
     J = np.random.randint(20, size=(100,))
 
+    return A, R, X, I, J
+
+
+def test_opsax(valid_arguments):
+    A, R, X, I, J = valid_arguments
+
+    reference = ref_opsax(A, R, X, I, J)
+    actual = opsax(A, R, X, I, J)
+    assert np.allclose(reference, actual)
+
+
+def test_opsax_no_neighbors(valid_arguments):
+    A, R, X, I, J = valid_arguments
+    # substitute all 1s by 2s so as to test the no-neighbor case
+    I[I == 1] = 2
+
     reference = ref_opsax(A, R, X, I, J)
     actual = opsax(A, R, X, I, J)
     assert np.allclose(reference, actual)
 
 
-def test_opsax_wrong_type():
-    with pytest.raises(ValueError):
-        opsax(np.array([1]), 2, 3, 4, 5)
+def test_opsax_wrong_type(valid_arguments):
+    A, R, X, I, J = valid_arguments
+    A = A.astype(np.int32)
+
+    with pytest.raises(TypeError, match="Wrong dtype for A in opsax: got int32"):
+        opsax(A, R, X, I, J)
+
+
+def test_opsax_wrong_number_of_dimensions(valid_arguments):
+    A, R, X, I, J = valid_arguments
+    A = A[..., np.newaxis]
+
+    with pytest.raises(
+        ValueError, match="A must be a 2D array in opsax, got a 3D array"
+    ):
+        opsax(A, R, X, I, J)
+
+
+def test_opsax_size_mismatch(valid_arguments):
+    A, R, X, I, J = valid_arguments
+    I = I[:5]
+
+    with pytest.raises(
+        mops.status.MopsError,
+        match="Dimension mismatch: the sizes of A along "
+        "dimension 0 and I along dimension 0 must match in opsax",
+    ):
+        opsax(A, R, X, I, J)
+
+
+def test_opsax_out_of_bounds(valid_arguments):
+    A, R, X, I, J = valid_arguments
+    I[0] = X.shape[0]
+
+    with pytest.raises(
+        mops.status.MopsError,
+        match="Index array I in operation opsax contains elements up to 20; "
+        "this would cause out-of-bounds accesses. With the provided "
+        "parameters, it can only contain elements up to 19",
+    ):
+        opsax(A, R, X, I, J)

python/mops/tests/sap.py

@@ -2,26 +2,68 @@
 import pytest
 from mops.reference_implementations import sparse_accumulation_of_products as ref_sap
 
+import mops
 from mops import sparse_accumulation_of_products as sap
 
 np.random.seed(0xDEADBEEF)
 
 
-def test_sap():
+@pytest.fixture
+def valid_arguments():
     A = np.random.rand(100, 20)
     B = np.random.rand(100, 6)
     C = np.random.rand(30)
-
     P_A = np.random.randint(20, size=(30,))
     P_B = np.random.randint(6, size=(30,))
     n_O = 35
     P_O = np.random.randint(n_O, size=(30,))
+    return A, B, C, P_A, P_B, P_O, n_O
+
+
+def test_sap(valid_arguments):
+    A, B, C, P_A, P_B, P_O, n_O = valid_arguments
 
     reference = ref_sap(A, B, C, P_A, P_B, P_O, n_O)
     actual = sap(A, B, C, P_A, P_B, P_O, n_O)
     assert np.allclose(reference, actual)
 
 
-def test_sap_wrong_type():
-    with pytest.raises(ValueError):
-        sap(np.array(1), 2, 3, 4, 5, 6, 7)
+def test_sap_wrong_type(valid_arguments):
+    A, B, C, P_A, P_B, P_O, n_O = valid_arguments
+    A = A.astype(np.int32)
+
+    with pytest.raises(TypeError, match="Wrong dtype for A in sap: got int32"):
+        sap(A, B, C, P_A, P_B, P_O, n_O)
+
+
+def test_sap_wrong_number_of_dimensions(valid_arguments):
+    A, B, C, P_A, P_B, P_O, n_O = valid_arguments
+    A = A[..., np.newaxis]
+
+    with pytest.raises(ValueError, match="A must be a 2D array in sap, got a 3D array"):
+        sap(A, B, C, P_A, P_B, P_O, n_O)
+
+
+def test_sap_size_mismatch(valid_arguments):
+    A, B, C, P_A, P_B, P_O, n_O = valid_arguments
+    C = C[:5]
+
+    with pytest.raises(
+        mops.status.MopsError,
+        match="Dimension mismatch: the sizes of C along "
+        "dimension 0 and P_A along dimension 0 must match in sap",
+    ):
+        sap(A, B, C, P_A, P_B, P_O, n_O)
+
+
+def test_sap_out_of_bounds(valid_arguments):
+    A, B, C, P_A, P_B, P_O, n_O = valid_arguments
+    P_A[0] = A.shape[1]
+
+    with pytest.raises(
+        mops.status.MopsError,
+        match="Index array P_A in operation sap contains elements up to 20; "
+        "this would cause out-of-bounds accesses. With the provided "
+        "parameters, it can only contain elements up to 19",
+    ):
+        sap(A, B, C, P_A, P_B, P_O, n_O)

python/mops/tests/sasax.py

@@ -4,12 +4,14 @@
     sparse_accumulation_scatter_add_with_weights as ref_sasax,
 )
 
+import mops
 from mops import sparse_accumulation_scatter_add_with_weights as sasax
 
 np.random.seed(0xDEADBEEF)
 
 
-def test_sasax():
+@pytest.fixture
+def valid_arguments():
     A = np.random.rand(100, 20)
     R = np.random.rand(100, 200)
     X = np.random.rand(25, 13, 200)
@@ -21,11 +23,65 @@ def test_sasax():
     M_2 = np.random.randint(13, size=(50,))
     M_3 = np.random.randint(n_O, size=(50,))
 
+    return A, R, X, C, I, J, M_1, M_2, M_3, n_O
+
+
+def test_sasax(valid_arguments):
+    A, R, X, C, I, J, M_1, M_2, M_3, n_O = valid_arguments
+
+    reference = ref_sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O)
+    actual = sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O)
+    assert np.allclose(reference, actual)
+
+
+def test_sasax_no_neighbors(valid_arguments):
+    A, R, X, C, I, J, M_1, M_2, M_3, n_O = valid_arguments
+    # substitute all 1s by 2s so as to test the no-neighbor case
+    I[I == 1] = 2
+
     reference = ref_sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O)
     actual = sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O)
     assert np.allclose(reference, actual)
 
 
-def test_sasax_wrong_type():
-    with pytest.raises(ValueError):
-        sasax(np.array([1]), 2, 3, 4, 5, 6, 7, 8, 9, 10)
+def test_sasax_wrong_type(valid_arguments):
+    A, R, X, C, I, J, M_1, M_2, M_3, n_O = valid_arguments
+    A = A.astype(np.int32)
+
+    with pytest.raises(TypeError, match="Wrong dtype for A in sasax: got int32"):
+        sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O)
+
+
+def test_sasax_wrong_number_of_dimensions(valid_arguments):
+    A, R, X, C, I, J, M_1, M_2, M_3, n_O = valid_arguments
+    A = A[..., np.newaxis]
+
+    with pytest.raises(
+        ValueError, match="A must be a 2D array in sasax, got a 3D array"
+    ):
+        sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O)
+
+
+def test_sasax_size_mismatch(valid_arguments):
+    A, R, X, C, I, J, M_1, M_2, M_3, n_O = valid_arguments
+    I = I[:5]
+
+    with pytest.raises(
+        mops.status.MopsError,
+        match="Dimension mismatch: the sizes of A along "
+        "dimension 0 and I along dimension 0 must match in sasax",
+    ):
+        sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O)
+
+
+def test_sasax_out_of_bounds(valid_arguments):
+    A, R, X, C, I, J, M_1, M_2, M_3, n_O = valid_arguments
+    I[0] = X.shape[0]
+
+    with pytest.raises(
+        mops.status.MopsError,
+        match="Index array I in operation sasax contains elements up to 25; "
+        "this would cause out-of-bounds accesses. With the provided "
+        "parameters, it can only contain elements up to 24",
+    ):
+        sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O)