Update and add more tests for solve()

Author: vlad-perevezentsev

dpnp/tests/test_linalg.py

@@ -2694,6 +2694,36 @@ def test_solve(self, dtype):
 
         assert_allclose(expected, result, rtol=1e-06)
 
+    @testing.with_requires("numpy>=2.0")
+    @pytest.mark.parametrize("dtype", get_float_complex_dtypes())
+    @pytest.mark.parametrize(
+        "a_shape, b_shape",
+        [
+            ((4, 4), (2, 2, 4, 3)),
+            ((2, 5, 5), (1, 5, 3)),
+            ((2, 4, 4), (2, 2, 4, 2)),
+            ((3, 2, 2), (3, 1, 2, 1)),
+            ((2, 2, 2, 2, 2), (2,)),
+            ((2, 2, 2, 2, 2), (2, 3)),
+        ],
+    )
+    def test_solve_broadcast(self, a_shape, b_shape, dtype):
+        # Set seed_value=81 to prevent
+        # random generation of the input singular matrix
+        a_np = generate_random_numpy_array(a_shape, dtype, seed_value=81)
+
+        # Set seed_value=76 to prevent
+        # random generation of the input singular matrix
+        b_np = generate_random_numpy_array(b_shape, dtype, seed_value=76)
+
+        a_dp = inp.array(a_np)
+        b_dp = inp.array(b_np)
+
+        expected = numpy.linalg.solve(a_np, b_np)
+        result = inp.linalg.solve(a_dp, b_dp)
+
+        assert_dtype_allclose(result, expected)
+
     @pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True))
     def test_solve_nrhs_greater_n(self, dtype):
         # Test checking the case when nrhs > n for

dpnp/tests/third_party/cupy/linalg_tests/test_solve.py

@@ -64,6 +64,9 @@ def test_solve(self):
             self.check_x((2, 4, 4), (2, 4))
             self.check_x((2, 3, 2, 2), (2, 3, 2))
             self.check_x((0, 2, 2), (0, 2))
+        else:  # Allowed since numpy 2.0
+            self.check_x((2, 3, 3), (3,))
+            self.check_x((2, 5, 3, 3), (3,))
 
     def check_shape(self, a_shape, b_shape, error_types):
         for xp, error_type in error_types.items():
@@ -96,11 +99,38 @@ def test_invalid_shape(self):
         self.check_shape((3, 3), (2,), value_errors)
         self.check_shape((3, 3), (2, 2), value_errors)
         self.check_shape((3, 3, 4), (3,), linalg_errors)
+        self.check_shape((3, 3), (0,), value_errors)
+        self.check_shape((0, 3, 4), (3,), linalg_errors)
         # Since numpy >= 2.0, this case does not raise an error
         if numpy.lib.NumpyVersion(numpy.__version__) < "2.0.0":
             self.check_shape((2, 3, 3), (3,), value_errors)
-        self.check_shape((3, 3), (0,), value_errors)
-        self.check_shape((0, 3, 4), (3,), linalg_errors)
+        else:
+            # Not allowed since numpy 2
+            self.check_shape(
+                (0, 2, 2),
+                (
+                    0,
+                    2,
+                ),
+                value_errors,
+            )
+            self.check_shape(
+                (2, 4, 4),
+                (
+                    2,
+                    4,
+                ),
+                value_errors,
+            )
+            self.check_shape(
+                (2, 3, 2, 2),
+                (
+                    2,
+                    3,
+                    2,
+                ),
+                value_errors,
+            )
 
 
 @testing.parameterize(