diff --git a/python/paddle/fluid/tests/unittests/test_linalg_lstsq_op.py b/python/paddle/fluid/tests/unittests/test_linalg_lstsq_op.py index 82576ab1bd1bf7dd0c5ac9deefebda4b028eab20..94dc901a56d0c92d47dd95f4fe1029e4919a2571 100644 --- a/python/paddle/fluid/tests/unittests/test_linalg_lstsq_op.py +++ b/python/paddle/fluid/tests/unittests/test_linalg_lstsq_op.py @@ -278,5 +278,38 @@ class LinalgLstsqTestCaseLarge2(LinalgLstsqTestCase): self._input_shape_2 = (50, 300) +class TestLinalgLstsqAPIError(unittest.TestCase): + def setUp(self): + pass + + def test_api_errors(self): + def test_x_bad_shape(): + x = paddle.to_tensor(np.random.random(size=(5)), dtype=np.float32) + y = paddle.to_tensor( + np.random.random(size=(5, 15)), dtype=np.float32 + ) + out = paddle.linalg.lstsq(x, y, driver='gelsy') + + def test_y_bad_shape(): + x = paddle.to_tensor( + np.random.random(size=(5, 10)), dtype=np.float32 + ) + y = paddle.to_tensor(np.random.random(size=(5)), dtype=np.float32) + out = paddle.linalg.lstsq(x, y, driver='gelsy') + + def test_shape_dismatch(): + x = paddle.to_tensor( + np.random.random(size=(5, 10)), dtype=np.float32 + ) + y = paddle.to_tensor( + np.random.random(size=(4, 15)), dtype=np.float32 + ) + out = paddle.linalg.lstsq(x, y, driver='gelsy') + + self.assertRaises(ValueError, test_x_bad_shape) + self.assertRaises(ValueError, test_y_bad_shape) + self.assertRaises(ValueError, test_shape_dismatch) + + if __name__ == '__main__': unittest.main() diff --git a/python/paddle/tensor/linalg.py b/python/paddle/tensor/linalg.py index 4cce1b01968a196250d9346ec4a4e173e21f6892..46f11130c035449480d05564c8055dc5c9d4b340 100644 --- a/python/paddle/tensor/linalg.py +++ b/python/paddle/tensor/linalg.py @@ -3171,13 +3171,26 @@ def lstsq(x, y, rcond=None, driver=None, name=None): else: raise RuntimeError("Only support lstsq api for CPU or CUDA device.") - if x.dtype == y.dtype and x.dtype in (paddle.float32, paddle.float64): - pass - else: + if not (x.dtype == y.dtype and x.dtype in (paddle.float32, paddle.float64)): raise ValueError( "Only support x and y have the same dtype such as 'float32' and 'float64'." ) + if x.ndim < 2: + raise ValueError( + f"The shape of x should be (*, M, N), but received ndim is [{x.ndim} < 2]" + ) + + if y.ndim < 2: + raise ValueError( + f"The shape of y should be (*, M, K), but received ndim is [{y.ndim} < 2]" + ) + + if x.shape[-2] != y.shape[-2]: + raise ValueError( + f"x with shape (*, M = {x.shape[-2]}, N) and y with shape (*, M = {y.shape[-2]}, K) should have same M." + ) + if rcond is None: if x.dtype == paddle.float32: rcond = 1e-7 * max(x.shape[-2], x.shape[-1])