Apply review comments

Author: mtsokol

ci/test_examples.sh

@@ -1,3 +1,3 @@
 for example in $(find ./examples/ -iname *.py); do
-  CI_MODE=True python $example
+  CI_MODE=1 python $example
 done

ci/test_notebooks.sh

@@ -1 +1 @@
-CI_MODE=True pytest -n 4 --nbmake --nbmake-timeout=600 ./examples/*.ipynb
+CI_MODE=1 pytest -n 4 --nbmake --nbmake-timeout=600 ./examples/*.ipynb

examples/elemwise_example.py

@@ -34,10 +34,9 @@
         s2 = sparse.asarray(s2_sps.asformat("csc"), format="csc")
 
         func = getattr(sparse, func_name)
-        # Compile
-        result_finch = func(s1, s2)
-        # Benchmark
-        benchmark(func, args=[s1, s2], info="Finch", iters=ITERS)
+
+        # Compile & Benchmark
+        result_finch = benchmark(func, args=[s1, s2], info="Finch", iters=ITERS)
 
         # ======= Numba =======
         os.environ[sparse._ENV_VAR_NAME] = "Numba"
@@ -47,10 +46,9 @@
         s2 = sparse.asarray(s2_sps)
 
         func = getattr(sparse, func_name)
-        # Compile
-        result_numba = func(s1, s2)
-        # Benchmark
-        benchmark(func, args=[s1, s2], info="Numba", iters=ITERS)
+
+        # Compile & Benchmark
+        result_numba = benchmark(func, args=[s1, s2], info="Numba", iters=ITERS)
 
         # ======= SciPy =======
         s1 = s1_sps
@@ -63,9 +61,8 @@
         elif func_name == "greater_equal":
             func, args = operator.ge, [s1, s2]
 
-        result_scipy = func(*args)
-        # Benchmark
-        benchmark(func, args=args, info="SciPy", iters=ITERS)
+        # Compile & Benchmark
+        result_scipy = benchmark(func, args=args, info="SciPy", iters=ITERS)
 
         np.testing.assert_allclose(result_numba.todense(), result_scipy.toarray())
         np.testing.assert_allclose(result_finch.todense(), result_numba.todense())

examples/matmul_example.py

@@ -33,10 +33,8 @@
     def sddmm_finch(a, b):
         return a @ b
 
-    # Compile
-    result_finch = sddmm_finch(a, b)
-    # Benchmark
-    benchmark(sddmm_finch, args=[a, b], info="Finch", iters=ITERS)
+    # Compile & Benchmark
+    result_finch = benchmark(sddmm_finch, args=[a, b], info="Finch", iters=ITERS)
 
     # ======= Numba =======
     os.environ[sparse._ENV_VAR_NAME] = "Numba"
@@ -48,10 +46,8 @@ def sddmm_finch(a, b):
     def sddmm_numba(a, b):
         return a @ b
 
-    # Compile
-    result_numba = sddmm_numba(a, b)
-    # Benchmark
-    benchmark(sddmm_numba, args=[a, b], info="Numba", iters=ITERS)
+    # Compile & Benchmark
+    result_numba = benchmark(sddmm_numba, args=[a, b], info="Numba", iters=ITERS)
 
     # ======= SciPy =======
     def sddmm_scipy(a, b):
@@ -60,9 +56,8 @@ def sddmm_scipy(a, b):
     a = a_sps
     b = b_sps
 
-    result_scipy = sddmm_scipy(a, b)
-    # Benchmark
-    benchmark(sddmm_scipy, args=[a, b], info="SciPy", iters=ITERS)
+    # Compile & Benchmark
+    result_scipy = benchmark(sddmm_scipy, args=[a, b], info="SciPy", iters=ITERS)
 
     # np.testing.assert_allclose(result_numba.todense(), result_scipy.toarray())
     # np.testing.assert_allclose(result_finch.todense(), result_numba.todense())

examples/mttkrp_example.py

@@ -35,10 +35,8 @@
     def mttkrp_finch(B, D, C):
         return sparse.sum(B[:, :, :, None] * D[None, None, :, :] * C[None, :, None, :], axis=(1, 2))
 
-    # Compile
-    result_finch = mttkrp_finch(B, D, C)
-    # Benchmark
-    benchmark(mttkrp_finch, args=[B, D, C], info="Finch", iters=ITERS)
+    # Compile & Benchmark
+    result_finch = benchmark(mttkrp_finch, args=[B, D, C], info="Finch", iters=ITERS)
 
     # ======= Numba =======
     os.environ[sparse._ENV_VAR_NAME] = "Numba"
@@ -51,9 +49,7 @@ def mttkrp_finch(B, D, C):
     def mttkrp_numba(B, D, C):
         return sparse.sum(B[:, :, :, None] * D[None, None, :, :] * C[None, :, None, :], axis=(1, 2))
 
-    # Compile
-    result_numba = mttkrp_numba(B, D, C)
-    # Benchmark
-    benchmark(mttkrp_numba, args=[B, D, C], info="Numba", iters=ITERS)
+    # Compile & Benchmark
+    result_numba = benchmark(mttkrp_numba, args=[B, D, C], info="Numba", iters=ITERS)
 
     np.testing.assert_allclose(result_finch.todense(), result_numba.todense())

examples/sddmm_example.py

@@ -37,10 +37,8 @@ def sddmm_finch(s, a, b):
             axis=-1,
         )
 
-    # Compile
-    result_finch = sddmm_finch(s, a, b)
-    # Benchmark
-    benchmark(sddmm_finch, args=[s, a, b], info="Finch", iters=ITERS)
+    # Compile & Benchmark
+    result_finch = benchmark(sddmm_finch, args=[s, a, b], info="Finch", iters=ITERS)
 
     # ======= Numba =======
     os.environ[sparse._ENV_VAR_NAME] = "Numba"
@@ -53,10 +51,8 @@ def sddmm_finch(s, a, b):
     def sddmm_numba(s, a, b):
         return s * (a @ b)
 
-    # Compile
-    result_numba = sddmm_numba(s, a, b)
-    # Benchmark
-    benchmark(sddmm_numba, args=[s, a, b], info="Numba", iters=ITERS)
+    # Compile & Benchmark
+    result_numba = benchmark(sddmm_numba, args=[s, a, b], info="Numba", iters=ITERS)
 
     # ======= SciPy =======
     def sddmm_scipy(s, a, b):
@@ -66,9 +62,8 @@ def sddmm_scipy(s, a, b):
     a = a_sps
     b = b_sps
 
-    result_scipy = sddmm_scipy(s, a, b)
-    # Benchmark
-    benchmark(sddmm_scipy, args=[s, a, b], info="SciPy", iters=ITERS)
+    # Compile & Benchmark
+    result_scipy = benchmark(sddmm_scipy, args=[s, a, b], info="SciPy", iters=ITERS)
 
     np.testing.assert_allclose(result_numba.todense(), result_scipy.toarray())
     np.testing.assert_allclose(result_finch.todense(), result_numba.todense())

examples/spmv_add_example.py

@@ -33,10 +33,8 @@
     def spmv_finch(A, x, y):
         return sparse.sum(A[:, None, :] * sparse.permute_dims(x, (1, 0))[None, :, :], axis=-1) + y
 
-    # Compile
-    result_finch = spmv_finch(A, x, y)
-    # Benchmark
-    benchmark(spmv_finch, args=[A, x, y], info="Finch", iters=ITERS)
+    # Compile & Benchmark
+    result_finch = benchmark(spmv_finch, args=[A, x, y], info="Finch", iters=ITERS)
 
     # ======= Numba =======
     os.environ[sparse._ENV_VAR_NAME] = "Numba"
@@ -49,11 +47,8 @@ def spmv_finch(A, x, y):
     def spmv_numba(A, x, y):
         return A @ x + y
 
-    # Compile
-    result_numba = spmv_numba(A, x, y)
-    assert sparse.nonzero(result_numba)[0].size > 5
-    # Benchmark
-    benchmark(spmv_numba, args=[A, x, y], info="Numba", iters=ITERS)
+    # Compile & Benchmark
+    result_numba = benchmark(spmv_numba, args=[A, x, y], info="Numba", iters=ITERS)
 
     # ======= SciPy =======
     def spmv_scipy(A, x, y):
@@ -63,9 +58,8 @@ def spmv_scipy(A, x, y):
     x = x_sps
     y = y_sps
 
-    result_scipy = spmv_scipy(A, x, y)
-    # Benchmark
-    benchmark(spmv_scipy, args=[A, x, y], info="SciPy", iters=ITERS)
+    # Compile & Benchmark
+    result_scipy = benchmark(spmv_scipy, args=[A, x, y], info="SciPy", iters=ITERS)
 
     np.testing.assert_allclose(result_numba, result_scipy)
     np.testing.assert_allclose(result_finch.todense(), result_numba)

examples/utils.py

@@ -3,17 +3,28 @@
 from collections.abc import Callable, Iterable
 from typing import Any
 
-CI_MODE = os.getenv("CI_MODE", default=False)
+CI_MODE = bool(int(os.getenv("CI_MODE", default="0")))
 
 
-def benchmark(func: Callable, args: Iterable[Any], info: str, iters: int):
+def benchmark(
+    func: Callable,
+    args: Iterable[Any],
+    info: str,
+    iters: int,
+) -> object:
+    # Compile
+    result = func(*args)
+
     if CI_MODE:
         print("CI mode - skipping benchmark")
-        return
+        return result
 
+    # Benchmark
     print(info)
     start = time.time()
     for _ in range(iters):
         func(*args)
     elapsed = time.time() - start
     print(f"Took {elapsed / iters} s.\n")
+
+    return result