Add tests for tensorsolve

Author: vlad-perevezentsev

tests/test_linalg.py

@@ -1409,3 +1409,24 @@ def test_pinv_errors(self):
         a_dp_q = inp.array(a_dp, sycl_queue=a_queue)
         rcond_dp_q = inp.array([0.5], dtype="float32", sycl_queue=rcond_queue)
         assert_raises(ValueError, inp.linalg.pinv, a_dp_q, rcond_dp_q)
+
+
+class TestTensorsolve:
+    def test_tensorsolve_errors(self):
+        a_dp = inp.eye(24, dtype="float32").reshape(4, 6, 8, 3)
+        b_dp = inp.ones(a_dp.shape[:2], dtype="float32")
+
+        # unsupported type `a` and `b`
+        a_np = inp.asnumpy(a_dp)
+        b_np = inp.asnumpy(b_dp)
+        assert_raises(TypeError, inp.linalg.tensorsolve, a_np, b_dp)
+        assert_raises(TypeError, inp.linalg.tensorsolve, a_dp, b_np)
+
+        # unsupported type `axes`
+        assert_raises(TypeError, inp.linalg.tensorsolve, a_dp, 2.0)
+        assert_raises(TypeError, inp.linalg.tensorsolve, a_dp, -2)
+
+        # incorrect axes
+        assert_raises(
+            inp.linalg.LinAlgError, inp.linalg.tensorsolve, a_dp, b_dp, (1,)
+        )

tests/test_sycl_queue.py

@@ -1849,3 +1849,24 @@ def test_pinv(shape, hermitian, rcond_as_array, device):
     B_queue = B_result.sycl_queue
 
     assert_sycl_queue_equal(B_queue, a_dp.sycl_queue)
+
+
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_tensorsolve(device):
+    a_np = numpy.random.randn(3, 2, 6).astype(dpnp.default_float_type())
+    b_np = numpy.ones(a_np.shape[:2], dtype=a_np.dtype)
+
+    a_dp = dpnp.array(a_np, device=device)
+    b_dp = dpnp.array(b_np, device=device)
+
+    result = dpnp.linalg.tensorsolve(a_dp, b_dp)
+    expected = numpy.linalg.tensorsolve(a_np, b_np)
+    assert_dtype_allclose(result, expected)
+
+    result_queue = result.sycl_queue
+
+    assert_sycl_queue_equal(result_queue, a_dp.sycl_queue)

tests/test_usm_type.py

@@ -1014,3 +1014,17 @@ def test_qr(shape, mode, usm_type):
 
         assert a.usm_type == dp_q.usm_type
         assert a.usm_type == dp_r.usm_type
+
+
+@pytest.mark.parametrize("usm_type_a", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize("usm_type_b", list_of_usm_types, ids=list_of_usm_types)
+def test_tensorsolve(usm_type_a, usm_type_b):
+    data = numpy.random.randn(3, 2, 6)
+    a = dp.array(data, usm_type=usm_type_a)
+    b = dp.ones(a.shape[:2], dtype=a.dtype, usm_type=usm_type_b)
+
+    result = dp.linalg.tensorsolve(a, b)
+
+    assert a.usm_type == usm_type_a
+    assert b.usm_type == usm_type_b
+    assert result.usm_type == du.get_coerced_usm_type([usm_type_a, usm_type_b])

tests/third_party/cupy/linalg_tests/test_solve.py

@@ -95,6 +95,26 @@ def test_invalid_shape(self):
         self.check_shape((0, 3, 4), (3,), linalg_errors)
 
 
+@testing.parameterize(
+    *testing.product(
+        {
+            "a_shape": [(2, 3, 6), (3, 4, 4, 3)],
+            "axes": [None, (0, 2)],
+        }
+    )
+)
+@testing.fix_random()
+class TestTensorSolve(unittest.TestCase):
+    @testing.for_dtypes("ifdFD")
+    @testing.numpy_cupy_allclose(atol=0.02, type_check=has_support_aspect64())
+    def test_tensorsolve(self, xp, dtype):
+        a_shape = self.a_shape
+        b_shape = self.a_shape[:2]
+        a = testing.shaped_random(a_shape, xp, dtype=dtype, seed=0)
+        b = testing.shaped_random(b_shape, xp, dtype=dtype, seed=1)
+        return xp.linalg.tensorsolve(a, b, axes=self.axes)
+
+
 @testing.parameterize(
     *testing.product(
         {