Merge 5157356 into 787e43b

Author: vlad-perevezentsev

dpnp/linalg/dpnp_iface_linalg.py

@@ -1612,12 +1612,12 @@ def solve(a, b):
     ----------
     a : (..., M, M) {dpnp.ndarray, usm_ndarray}
         Coefficient matrix.
-    b : {(…, M,), (…, M, K)} {dpnp.ndarray, usm_ndarray}
+    b : {(M,), (..., M, K)} {dpnp.ndarray, usm_ndarray}
         Ordinate or "dependent variable" values.
 
     Returns
     -------
-    out : {(…, M,), (…, M, K)} dpnp.ndarray
+    out : {(..., M,), (..., M, K)} dpnp.ndarray
         Solution to the system `ax = b`. Returned shape is identical to `b`.
 
     See Also
@@ -1644,14 +1644,53 @@ def solve(a, b):
     assert_stacked_2d(a)
     assert_stacked_square(a)
 
-    if not (
-        a.ndim in [b.ndim, b.ndim + 1] and a.shape[:-1] == b.shape[: a.ndim - 1]
-    ):
-        raise dpnp.linalg.LinAlgError(
-            "a must have (..., M, M) shape and b must have (..., M) "
-            "or (..., M, K)"
+    a_ndim = a.ndim
+    b_ndim = b.ndim
+
+    a_shape = a.shape
+    b_shape = b.shape
+
+    if numpy.lib.NumpyVersion(numpy.__version__) < "2.0.0":
+        if not (
+            a_ndim in [b_ndim, b_ndim + 1]
+            and a_shape[:-1] == b_shape[: a_ndim - 1]
+        ):
+            raise dpnp.linalg.LinAlgError(
+                "a must have (..., M, M) shape and b must have (..., M) "
+                "or (..., M, K)"
+            )
+
+    else:  # compatible with numpy>=2.0
+        if b_ndim == 0:
+            raise ValueError("b must have at least one dimension")
+        if b_ndim == 1:
+            if a_shape[-1] != b.size:
+                raise ValueError(
+                    "a must have (..., M, M) shape and b must have (M,) "
+                    "for one-dimensional b"
+                )
+            b = dpnp.broadcast_to(b, a_shape[:-1])
+            return dpnp_solve(a, b)
+
+        if a_shape[-1] != b_shape[-2]:
+            raise ValueError(
+                "a must have (..., M, M) shape and b must have "
+                "(..., M, K) shape"
+            )
+
+        # Use dpnp.broadcast_shapes() to align the resulting batch shapes
+        broadcasted_batch_shape = dpnp.broadcast_shapes(
+            a_shape[:-2], b_shape[:-2]
         )
 
+        a_broadcasted_shape = broadcasted_batch_shape + a_shape[-2:]
+        b_broadcasted_shape = broadcasted_batch_shape + b_shape[-2:]
+
+        if a_shape != a_broadcasted_shape:
+            a = dpnp.broadcast_to(a, a_broadcasted_shape)
+        if b_shape != b_broadcasted_shape:
+            b = dpnp.broadcast_to(b, b_broadcasted_shape)
+
     return dpnp_solve(a, b)
 
 

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,16 @@ 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(