Fix the solve() inputs strategy with the new, unambiguous input shapes

Author: asmeurer

array_api_tests/test_linalg.py

@@ -491,7 +491,7 @@ def solve_args():
     Strategy for the x1 and x2 arguments to test_solve()
 
     solve() takes x1, x2, where x1 is any stack of square invertible matrices
-    of shape (..., M, M), and x2 is either shape (..., M) or (..., M, K),
+    of shape (..., M, M), and x2 is either shape (M,) or (..., M, K),
     where the ... parts of x1 and x2 are broadcast compatible.
     """
     stack_shapes = shared(two_mutually_broadcastable_shapes)
@@ -501,26 +501,18 @@ def solve_args():
                                                                   pair[0])))
 
     @composite
-    def x2_shapes(draw):
-        end = draw(xps.array_shapes(min_dims=0, max_dims=1, min_side=0,
-                                    max_side=SQRT_MAX_ARRAY_SIZE))
-        return draw(stack_shapes)[1] + draw(x1).shape[-1:] + end
+    def _x2_shapes(draw):
+        end = draw(integers(0, SQRT_MAX_ARRAY_SIZE))
+        return draw(stack_shapes)[1] + draw(x1).shape[-1:] + (end,)
 
-    x2 = xps.arrays(dtype=xps.floating_dtypes(), shape=x2_shapes())
+    x2_shapes = one_of(x1.map(lambda x: (x.shape[-1],)), _x2_shapes())
+    x2 = xps.arrays(dtype=xps.floating_dtypes(), shape=x2_shapes)
     return x1, x2
 
 @pytest.mark.xp_extension('linalg')
 @given(*solve_args())
 def test_solve(x1, x2):
-    # TODO: solve() is currently ambiguous, in that some inputs can be
-    # interpreted in two different ways. For example, if x1 is shape (2, 2, 2)
-    # and x2 is shape (2, 2), should this be interpreted as x2 is (2,) stack
-    # of a (2,) vector, i.e., the result would be (2, 2, 2, 1) after
-    # broadcasting, or as a single stack of a 2x2 matrix, i.e., resulting in
-    # (2, 2, 2, 2).
-
-    # res = linalg.solve(x1, x2)
-    pass
+    res = linalg.solve(x1, x2)
 
 @pytest.mark.xp_extension('linalg')
 @given(