Add an rtols strategy

Start implementing tests for pinv and matrix_norm.

Author: asmeurer

array_api_tests/hypothesis_helpers.py

@@ -158,10 +158,21 @@ def matrix_shapes(draw, stack_shapes=shapes()):
 
 square_matrix_shapes = matrix_shapes().filter(lambda shape: shape[-1] == shape[-2])
 
-finite_matrices = xps.arrays(dtype=xps.floating_dtypes(),
-                             shape=matrix_shapes(),
-                             elements=dict(allow_nan=False,
-                                           allow_infinity=False))
+@composite
+def finite_matrices(draw, shape=matrix_shapes()):
+    return draw(xps.arrays(dtype=xps.floating_dtypes(),
+                           shape=shape,
+                           elements=dict(allow_nan=False,
+                                         allow_infinity=False)))
+
+rtol_shared_matrix_shapes = shared(matrix_shapes())
+# Should we set a max_value here?
+_rtol_float_kw = dict(allow_nan=False, allow_infinity=False, min_value=0)
+rtols = one_of(floats(**_rtol_float_kw),
+               xps.arrays(dtype=xps.floating_dtypes(),
+                          shape=rtol_shared_matrix_shapes.map(lambda shape:  shape[:-2]),
+                          elements=_rtol_float_kw))
+
 
 def mutually_broadcastable_shapes(
     num_shapes: int,

array_api_tests/test_linalg.py

@@ -26,7 +26,8 @@
                                  invertible_matrices, two_mutual_arrays,
                                  mutually_promotable_dtypes, one_d_shapes,
                                  two_mutually_broadcastable_shapes,
-                                 SQRT_MAX_ARRAY_SIZE, finite_matrices)
+                                 SQRT_MAX_ARRAY_SIZE, finite_matrices,
+                                 rtol_shared_matrix_shapes, rtols)
 from . import dtype_helpers as dh
 from . import pytest_helpers as ph
 
@@ -311,7 +312,7 @@ def test_matmul(x1, x2):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=finite_matrices,
+    x=finite_matrices(),
     kw=kwargs(keepdims=booleans(),
               ord=sampled_from([-float('inf'), -2, -2, 1, 2, float('inf'), 'fro', 'nuc']))
 )
@@ -357,12 +358,11 @@ def test_matrix_power(x, n):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=xps.arrays(dtype=xps.floating_dtypes(), shape=shapes()),
-    kw=kwargs(rtol=todo)
+    x=finite_matrices(shape=rtol_shared_matrix_shapes),
+    kw=kwargs(rtol=rtols)
 )
 def test_matrix_rank(x, kw):
-    # res = linalg.matrix_rank(x, **kw)
-    pass
+    res = linalg.matrix_rank(x, **kw)
 
 @given(
     x=xps.arrays(dtype=dtypes, shape=matrix_shapes()),
@@ -407,12 +407,11 @@ def test_outer(x1, x2):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=xps.arrays(dtype=xps.floating_dtypes(), shape=shapes()),
-    kw=kwargs(rtol=todo)
+    x=finite_matrices(shape=rtol_shared_matrix_shapes),
+    kw=kwargs(rtol=rtols)
 )
 def test_pinv(x, kw):
-    # res = linalg.pinv(x, **kw)
-    pass
+    res = linalg.pinv(x, **kw)
 
 @pytest.mark.xp_extension('linalg')
 @given(
@@ -525,7 +524,7 @@ def test_solve(x1, x2):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=finite_matrices,
+    x=finite_matrices(),
     kw=kwargs(full_matrices=booleans())
 )
 def test_svd(x, kw):
@@ -561,7 +560,7 @@ def test_svd(x, kw):
 
 @pytest.mark.xp_extension('linalg')
 @given(
-    x=finite_matrices,
+    x=finite_matrices(),
 )
 def test_svdvals(x):
     res = linalg.svdvals(x)