Overaul hypothesis strategies

Improve readability. 
Given infinite hypothesis examples, this commit does not introduce any functional changes.
However, given a fixed and relatively low max_examples setting, it substantially increases test coverage:

1. It prevents examples that were previously skipped by assume(). According to pytest --hypothesis-show-statistics, these were ~10% for dotv and ~20% for gemm.

2. It prevents examples that ended up being duplicates due to trimming. For example, in dotv, hypothesis could previously generate examples e.g. A=[1,2,3,4], B=[5,6] and A=[1,2], B=[5,6,7,8]; both would get trimmed to A=[1,2], B=[5,6].

3. In gemm, it removes an entire degree of freedom by removing unused variable out_col, which again would result in duplicate examples.

Author: crusaderky

tests/blis_tests_common.py

@@ -1,34 +1,25 @@
 # Copyright ExplosionAI GmbH, released under BSD.
-import numpy as np
-
-np.random.seed(0)
-
 from hypothesis import settings
-from hypothesis.strategies import integers, floats
+from hypothesis import strategies as st
 from hypothesis.extra.numpy import arrays
 
 # Increase this to run more thorough tests
-hypothesis_default_profile = settings.register_profile("default", max_examples=100)
-
-
-def lengths(lo, hi):
-    return integers(min_value=lo, max_value=hi)
-
-
-def ndarrays_of_shape(shape, lo, hi, dtype):
-    width = 64 if dtype == "float64" else 32
-    return arrays(
-        dtype,
-        shape=shape,
-        elements=floats(
-            min_value=lo,
-            max_value=hi,
-            width=width,
-        ),
-    )
-
-
-def ndarrays(min_len, max_len, min_val, max_val, dtype):
-    return lengths(lo=min_len, hi=max_len).flatmap(
-        lambda n: ndarrays_of_shape(n, lo=min_val, hi=max_val, dtype=dtype)
-    )
+hypothesis_default_profile = settings.register_profile("default", max_examples=200)
+
+# (dtype, atol, rtol)
+dtypes_tols = st.sampled_from(
+    [
+        # FIXME huge tolerance needed for float32:
+        # https://github.com/explosion/cython-blis/issues/142
+        ("float32", 1e-2, 1e-4),
+        ("float64", 1e-9, 1e-9),
+    ]
+)
+
+
+def ndarrays(shape, lo, hi, dtype):
+    """Draw ND NumPy arrays of floats"""
+    assert isinstance(dtype, str) and dtype.startswith("float")
+    width = int(dtype[5:])
+    elements = st.floats(min_value=lo, max_value=hi, width=width)
+    return arrays(dtype, shape=shape, elements=elements)

tests/test_dotv.py

@@ -1,45 +1,31 @@
 # Copyright ExplosionAI GmbH, released under BSD.
 import numpy as np
 import pytest
-from hypothesis import given, assume
-
+from hypothesis import given
+from hypothesis import strategies as st
 from numpy.testing import assert_allclose
-from blis_tests_common import ndarrays
+
+from blis_tests_common import dtypes_tols, ndarrays
 from blis.py import dotv
 
 
+@st.composite
+def dotv_arrays(draw):
+    """Draw two 1D NumPy arrays of the same size and dtype"""
+    size = draw(st.integers(min_value=1, max_value=100))
+    dtype, atol, rtol = draw(dtypes_tols)
+    arr_st = ndarrays((size,), lo=-100.0, hi=100.0, dtype=dtype)
+    return draw(arr_st), draw(arr_st), atol, rtol
+
+
 def test_incompatible_shape():
     with pytest.raises(ValueError):
         dotv(np.zeros(2), np.zeros(3))
 
 
-@given(
-    ndarrays(min_len=1, max_len=100, min_val=-100.0, max_val=100.0, dtype="float64"),
-    ndarrays(min_len=1, max_len=100, min_val=-100.0, max_val=100.0, dtype="float64"),
-)
-def test_memoryview_double_noconj(A, B):
-    if len(A) < len(B):
-        B = B[: len(A)]
-    else:
-        A = A[: len(B)]
-    assume(A is not None)
-    assume(B is not None)
-    numpy_result = A.dot(B)
-    result = dotv(A, B)
-    assert_allclose(result, numpy_result, atol=1e-9, rtol=1e-9)
-
-
-@given(
-    ndarrays(min_len=1, max_len=100, min_val=-100.0, max_val=100.0, dtype="float32"),
-    ndarrays(min_len=1, max_len=100, min_val=-100.0, max_val=100.0, dtype="float32"),
-)
-def test_memoryview_float_noconj(A, B):
-    if len(A) < len(B):
-        B = B[: len(A)]
-    else:
-        A = A[: len(B)]
-    assume(A is not None)
-    assume(B is not None)
+@given(dotv_arrays())
+def test_memoryview_noconj(arrays):
+    A, B, atol, rtol = arrays
     numpy_result = A.dot(B)
     result = dotv(A, B)
-    assert_allclose(result, numpy_result, atol=1e-2, rtol=1e-4)
+    assert_allclose(result, numpy_result, atol=atol, rtol=rtol)

tests/test_gemm.py

@@ -1,39 +1,25 @@
 # Copyright ExplosionAI GmbH, released under BSD.
-from math import sqrt, floor
-
-from hypothesis import given, assume
-from hypothesis.strategies import integers
-
 import numpy as np
-from numpy.testing import assert_allclose
 import pytest
+from hypothesis import given
+from hypothesis import strategies as st
+from numpy.testing import assert_allclose
 
-from blis_tests_common import ndarrays
+from blis_tests_common import dtypes_tols, ndarrays
 from blis.py import gemm
 
 
-def _stretch_matrix(data, m, n):
-    ratio = sqrt(len(data) / (m * n))
-    m = int(floor(m * ratio))
-    n = int(floor(n * ratio))
-    data = np.ascontiguousarray(data[: m * n], dtype=data.dtype)
-    return data.reshape((m, n)), m, n
-
-
-def _reshape_for_gemm(
-    A, B, a_rows, a_cols, out_cols, dtype, trans_a=False, trans_b=False
-):
-    A, a_rows, a_cols = _stretch_matrix(A, a_rows, a_cols)
-    if len(B) < a_cols or a_cols < 1:
-        return (None, None, None)
-    b_cols = int(floor(len(B) / a_cols))
-    B = np.ascontiguousarray(B.flatten()[: a_cols * b_cols], dtype=dtype)
-    B = B.reshape((a_cols, b_cols))
-    out_cols = B.shape[1]
-    C = np.zeros(shape=(A.shape[0], B.shape[1]), dtype=dtype)
-    if trans_a:
-        A = np.ascontiguousarray(A.T, dtype=dtype)
-    return A, B, C
+@st.composite
+def gemm_arrays(draw):
+    """Draw two NumPy arrays with shapes (l, m) and (m, n) of the same dtype"""
+    max_size = 100
+    l = draw(st.integers(min_value=1, max_value=max_size))  # noqa: E741
+    n = draw(st.integers(min_value=1, max_value=max_size))
+    m = draw(st.integers(min_value=1, max_value=max_size // max(l, n)))
+    dtype, atol, rtol = draw(dtypes_tols)
+    A = draw(ndarrays((l, m), lo=-100.0, hi=100.0, dtype=dtype))
+    B = draw(ndarrays((m, n), lo=-100.0, hi=100.0, dtype=dtype))
+    return A, B, atol, rtol
 
 
 def test_incompatible_shape():
@@ -47,41 +33,10 @@ def test_incompatible_shape():
         gemm(np.zeros((3, 2)), np.zeros((3, 2)), trans1=True, trans2=True)
 
 
-@given(
-    ndarrays(min_len=1, max_len=100, min_val=-100.0, max_val=100.0, dtype="float64"),
-    ndarrays(min_len=1, max_len=100, min_val=-100.0, max_val=100.0, dtype="float64"),
-    integers(min_value=2, max_value=1000),
-    integers(min_value=2, max_value=1000),
-    integers(min_value=2, max_value=1000),
-)
-def test_memoryview_double_notrans(A, B, a_rows, a_cols, out_cols):
-    A, B, C = _reshape_for_gemm(A, B, a_rows, a_cols, out_cols, "float64")
-    assume(A is not None)
-    assume(B is not None)
-    assume(C is not None)
-    assume(A.size >= 1)
-    assume(B.size >= 1)
-    assume(C.size >= 1)
-    gemm(A, B, out=C)
+@given(gemm_arrays())
+def test_memoryview_notrans(arrays):
+    A, B, atol, rtol = arrays
     numpy_result = A.dot(B)
-    assert_allclose(C, numpy_result, atol=1e-9, rtol=1e-9)
-
-
-@given(
-    ndarrays(min_len=1, max_len=100, min_val=-100.0, max_val=100.0, dtype="float32"),
-    ndarrays(min_len=1, max_len=100, min_val=-100.0, max_val=100.0, dtype="float32"),
-    integers(min_value=2, max_value=1000),
-    integers(min_value=2, max_value=1000),
-    integers(min_value=2, max_value=1000),
-)
-def test_memoryview_float_notrans(A, B, a_rows, a_cols, out_cols):
-    A, B, C = _reshape_for_gemm(A, B, a_rows, a_cols, out_cols, dtype="float32")
-    assume(A is not None)
-    assume(B is not None)
-    assume(C is not None)
-    assume(A.size >= 1)
-    assume(B.size >= 1)
-    assume(C.size >= 1)
+    C = np.zeros_like(numpy_result)  # (l, n)
     gemm(A, B, out=C)
-    numpy_result = A.dot(B)
-    assert_allclose(C, numpy_result, atol=1e-2, rtol=1e-4)
+    assert_allclose(C, numpy_result, atol=atol, rtol=rtol)