diff --git a/python/paddle/linalg.py b/python/paddle/linalg.py index 4c2d6c00b9f0d..9c49ba3393145 100644 --- a/python/paddle/linalg.py +++ b/python/paddle/linalg.py @@ -24,6 +24,7 @@ eigh, eigvals, eigvalsh, + householder_product, lstsq, lu, lu_unpack, @@ -53,6 +54,7 @@ 'matrix_rank', 'svd', 'qr', + 'householder_product', 'pca_lowrank', 'lu', 'lu_unpack', diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py index d1f29d049eeca..d2324dfb2ae8d 100644 --- a/python/paddle/tensor/__init__.py +++ b/python/paddle/tensor/__init__.py @@ -60,6 +60,7 @@ zeros, zeros_like, ) + from .einsum import einsum # noqa: F401 from .linalg import ( # noqa: F401 bincount, @@ -78,6 +79,7 @@ eigvals, eigvalsh, histogram, + householder_product, lstsq, lu, lu_unpack, @@ -435,6 +437,7 @@ 'mv', 'matrix_power', 'qr', + 'householder_product', 'pca_lowrank', 'eigvals', 'eigvalsh', diff --git a/python/paddle/tensor/linalg.py b/python/paddle/tensor/linalg.py index 091bde960bacb..212125825d9b1 100644 --- a/python/paddle/tensor/linalg.py +++ b/python/paddle/tensor/linalg.py @@ -3739,3 +3739,133 @@ def cdist( return paddle.linalg.norm( x[..., None, :] - y[..., None, :, :], p=p, axis=-1 ) + + +def householder_product(x, tau, name=None): + r""" + + Computes the first n columns of a product of Householder matrices. + + This function can get the vector :math:`\omega_{i}` from matrix `x` (m x n), the :math:`i-1` elements are zeros, and the i-th is `1`, the rest of the elements are from i-th column of `x`. + And with the vector `tau` can calculate the first n columns of a product of Householder matrices. + + :math:`H_i = I_m - \tau_i \omega_i \omega_i^H` + + Args: + x (Tensor): A tensor with shape (*, m, n) where * is zero or more batch dimensions. + tau (Tensor): A tensor with shape (*, k) where * is zero or more batch dimensions. + name (str, optional): For details, please refer to :ref:`api_guide_Name`. Generally, no setting is required. Default: None. + + Returns: + Tensor, the dtype is same as input tensor, the Q in QR decomposition. + + :math:`out = Q = H_1H_2H_3...H_k` + + Examples: + .. code-block:: python + + >>> import paddle + >>> x = paddle.to_tensor([[-1.1280, 0.9012, -0.0190], + ... [ 0.3699, 2.2133, -1.4792], + ... [ 0.0308, 0.3361, -3.1761], + ... [-0.0726, 0.8245, -0.3812]]) + >>> tau = paddle.to_tensor([1.7497, 1.1156, 1.7462]) + >>> Q = paddle.linalg.householder_product(x, tau) + >>> print(Q) + Tensor(shape=[4, 3], dtype=float32, place=Place(gpu:0), stop_gradient=True, + [[-0.74969995, -0.02181768, 0.31115776], + [-0.64721400, -0.12367040, -0.21738708], + [-0.05389076, -0.37562513, -0.84836429], + [ 0.12702821, -0.91822827, 0.36892807]]) + """ + + check_dtype( + x.dtype, + 'x', + [ + 'float32', + 'float64', + 'complex64', + 'complex128', + ], + 'householder_product', + ) + check_dtype( + tau.dtype, + 'tau', + [ + 'float32', + 'float64', + 'complex64', + 'complex128', + ], + 'householder_product', + ) + assert ( + x.dtype == tau.dtype + ), "The input x must have the same dtype with input tau.\n" + assert ( + len(x.shape) >= 2 + and len(tau.shape) >= 1 + and len(x.shape) == len(tau.shape) + 1 + ), ( + "The input x must have more than 2 dimensions, and input tau must have more than 1 dimension," + "and the dimension of x is 1 larger than the dimension of tau\n" + ) + assert ( + x.shape[-2] >= x.shape[-1] + ), "The rows of input x must be greater than or equal to the columns of input x.\n" + assert ( + x.shape[-1] >= tau.shape[-1] + ), "The last dim of x must be greater than tau.\n" + for idx, _ in enumerate(x.shape[:-2]): + assert ( + x.shape[idx] == tau.shape[idx] + ), "The input x must have the same batch dimensions with input tau.\n" + + def _householder_product(x, tau): + m, n = x.shape[-2:] + k = tau.shape[-1] + Q = paddle.eye(m).astype(x.dtype) + for i in range(min(k, n)): + w = x[i:, i] + if in_dynamic_mode(): + w[0] = 1 + else: + w = paddle.static.setitem(w, 0, 1) + w = w.reshape([-1, 1]) + if in_dynamic_mode(): + if x.dtype in [paddle.complex128, paddle.complex64]: + Q[:, i:] = Q[:, i:] - ( + Q[:, i:] @ w @ paddle.conj(w).T * tau[i] + ) + else: + Q[:, i:] = Q[:, i:] - (Q[:, i:] @ w @ w.T * tau[i]) + else: + Q = paddle.static.setitem( + Q, + (slice(None), slice(i, None)), + Q[:, i:] - (Q[:, i:] @ w @ w.T * tau[i]) + if x.dtype in [paddle.complex128, paddle.complex64] + else Q[:, i:] - (Q[:, i:] @ w @ w.T * tau[i]), + ) + return Q[:, :n] + + if len(x.shape) == 2: + return _householder_product(x, tau) + m, n = x.shape[-2:] + org_x_shape = x.shape + org_tau_shape = tau.shape + x = x.reshape((-1, org_x_shape[-2], org_x_shape[-1])) + tau = tau.reshape((-1, org_tau_shape[-1])) + n_batch = x.shape[0] + out = paddle.zeros([n_batch, m, n], dtype=x.dtype) + for i in range(n_batch): + if in_dynamic_mode(): + out[i] = _householder_product(x[i], tau[i]) + else: + out = paddle.static.setitem( + out, i, _householder_product(x[i], tau[i]) + ) + out = out.reshape(org_x_shape) + return out diff --git a/test/legacy_test/test_householder_product.py b/test/legacy_test/test_householder_product.py new file mode 100644 index 0000000000000..50f676b8ccd1f --- /dev/null +++ b/test/legacy_test/test_householder_product.py @@ -0,0 +1,241 @@ +# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np + +import paddle + + +def geqrf(x): + def _geqrf(x): + m, n = x.shape + tau = np.zeros([n, 1], dtype=x.dtype) + for i in range(min(n, m)): + alpha = x[i, i] + normx = np.linalg.norm(x[min(i + 1, m) :, i]) + beta = np.linalg.norm(x[i:, i]) + if x.dtype in [np.complex64, np.complex128]: + s = 1 if alpha < 0 else -1 + else: + alphar = x[i, i].real + s = 1 if alphar < 0 else -1 + u1 = alpha - s * beta + w = x[i:, i] / u1 + w[0] = 1 + x[i + 1 :, i] = w[1 : m - i + 1] + if normx == 0: + tau[i] = 0 + else: + tau[i] = -s * u1 / beta + x[i, i] = s * beta + w = w.reshape([-1, 1]) + if x.dtype in [np.complex64, np.complex128]: + x[i:, i + 1 :] = x[i:, i + 1 :] - (tau[i] * w) @ ( + np.conj(w).T @ x[i:, i + 1 :] + ) + else: + x[i:, i + 1 :] = x[i:, i + 1 :] - (tau[i] * w) @ ( + w.T @ x[i:, i + 1 :] + ) + return x, tau[: min(m, n)].reshape(-1) + + if len(x.shape) == 2: + return _geqrf(x) + m, n = x.shape[-2:] + org_x_shape = x.shape + x = x.reshape((-1, x.shape[-2], x.shape[-1])) + n_batch = x.shape[0] + out = np.zeros([n_batch, m, n], dtype=x.dtype) + taus = np.zeros([n_batch, min(m, n)], dtype=x.dtype) + org_taus_shape = list(org_x_shape[:-2]) + [min(m, n)] + for i in range(n_batch): + out[i], t = _geqrf(x[i]) + taus[i, :] = t.reshape(-1) + return out.reshape(org_x_shape), taus.reshape(org_taus_shape) + + +def ref_qr(x): + def _ref_qr(x): + q, _ = np.linalg.qr(x) + return q + + if len(x.shape) == 2: + return _ref_qr(x) + m, n = x.shape[-2:] + org_shape = x.shape + x = x.reshape((-1, x.shape[-2], x.shape[-1])) + n_batch = x.shape[0] + out = np.zeros([n_batch, m, n]) + for i in range(n_batch): + out[i] = _ref_qr(x[i]) + return out.reshape(org_shape) + + +class TestHouseholderProductAPI(unittest.TestCase): + def setUp(self): + self.init_input() + self.place = ( + paddle.CUDAPlace(0) + if paddle.is_compiled_with_cuda() + else paddle.CPUPlace() + ) + + def init_input(self): + self.x = np.array( + [ + [1, 2, 4], + [0, 0, 5], + [0, 3, 6], + ], + dtype=np.float32, + ) + + def test_static_api(self): + m, n = self.x.shape[-2:] + self._x = self.x.copy() + self.geqrf_x, self.tau = geqrf(self.x) + paddle.enable_static() + with paddle.static.program_guard(paddle.static.Program()): + x = paddle.static.data( + 'x', self.geqrf_x.shape, dtype=self.geqrf_x.dtype + ) + tau = paddle.static.data( + 'tau', self.tau.shape, dtype=self.tau.dtype + ) + out = paddle.linalg.householder_product(x, tau) + exe = paddle.static.Executor(self.place) + res = exe.run( + feed={'x': self.geqrf_x, 'tau': self.tau}, fetch_list=[out] + ) + out_ref = ref_qr(self._x) + np.testing.assert_allclose(out_ref, res[0], atol=1e-3) + + def test_dygraph_api(self): + m, n = self.x.shape[-2:] + self._x = self.x.copy() + self.geqrf_x, self.tau = geqrf(self.x) + paddle.disable_static(self.place) + x = paddle.to_tensor(self.geqrf_x) + tau = paddle.to_tensor(self.tau) + out = paddle.linalg.householder_product(x, tau) + out_ref = ref_qr(self._x) + np.testing.assert_allclose(out_ref, out.numpy(), atol=1e-3) + paddle.enable_static() + + def test_error(self): + pass + + +class TestHouseholderProductAPICase1(TestHouseholderProductAPI): + def init_input(self): + self.x = np.random.randn(4, 3).astype('float32') + + +class TestHouseholderProductAPICase2(TestHouseholderProductAPI): + def init_input(self): + self.x = np.random.randn(4, 3).astype('float64') + + +class TestHouseholderProductAPICase3(TestHouseholderProductAPI): + def init_input(self): + self.x = np.random.randn(10, 2).astype('float32') + + +class TestHouseholderProductAPICase4(TestHouseholderProductAPI): + def init_input(self): + self.x = np.random.randn(5, 4, 3).astype('float32') + + +class TestHouseholderProductAPICase5(TestHouseholderProductAPI): + def init_input(self): + self.x = np.random.randn(4, 3).astype('float32') + + +# complex dtype +class TestHouseholderProductAPICase6(TestHouseholderProductAPI): + def init_input(self): + self.x = np.random.randn(4, 3).astype('complex64') + + +class TestHouseholderProductAPICase7(TestHouseholderProductAPI): + def init_input(self): + self.x = np.random.randn(4, 3).astype('complex128') + + +class TestHouseholderProductAPI_batch_error(TestHouseholderProductAPI): + # shape "*" in x.shape:[*, m, n] and tau.shape:[*, k] must be the same, eg. * == [2, 2] in x, but * == [2, 3] in tau + def test_error(self): + with self.assertRaises(AssertionError): + x = paddle.randn([2, 2, 5, 4]) + tau = paddle.randn([2, 3, 4]) + out = paddle.linalg.householder_product(x, tau) + + +class TestHouseholderProductAPI_dim_error(TestHouseholderProductAPI): + # len(x.shape) must be greater(equal) than 2, len(tau.shape) must be greater(equal) than 1 + def test_error(self): + with self.assertRaises(AssertionError): + x = paddle.to_tensor( + [ + 3, + ], + dtype=paddle.float32, + ) + tau = paddle.to_tensor([], dtype=paddle.float32) + out = paddle.linalg.householder_product(x, tau) + + +class TestHouseholderProductAPI_type_error(TestHouseholderProductAPI): + # type of x and tau must be float32 or float64 + def test_error(self): + with self.assertRaises(TypeError): + x = paddle.randn([3, 2, 1], dtype=paddle.int32) + tau = paddle.randn([3, 1], dtype=paddle.int32) + out = paddle.linalg.householder_product(x, tau) + + +class TestHouseholderProductAPI_shape_dismatch_error(TestHouseholderProductAPI): + # len(x.shape) and len(tau.shape) + 1 must be equal + def test_error(self): + with self.assertRaises(AssertionError): + x = paddle.randn([3, 2, 1], dtype=paddle.float32) + tau = paddle.randn([6, 2, 4], dtype=paddle.float32) + out = paddle.linalg.householder_product(x, tau) + + +class TestHouseholderProductAPI_col_row_error(TestHouseholderProductAPI): + # row must be bigger than col in x + def test_error(self): + with self.assertRaises(AssertionError): + x = paddle.randn([3, 6], dtype=paddle.float32) + tau = paddle.randn([6], dtype=paddle.float32) + out = paddle.linalg.householder_product(x, tau) + + +class TestHouseholderProductAPI_n_greater_than_k_error( + TestHouseholderProductAPI +): + # A.shape:[*, m, n], tau.shape:[*, k], n must be greater than k + def test_error(self): + with self.assertRaises(AssertionError): + x = paddle.randn([6, 3], dtype=paddle.float32) + tau = paddle.randn([4], dtype=paddle.float32) + out = paddle.linalg.householder_product(x, tau) + + +if __name__ == '__main__': + paddle.enable_static() + unittest.main()