wehub-resource-sync
439 lines
13 KiB
Python
439 lines
13 KiB
Python
# Copyright (c) 2025 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 sys
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle import base
|
|
|
|
sys.path.append("..")
|
|
from op_test import OpTest, get_places
|
|
|
|
|
|
def _transpose_last_2dim(x):
|
|
"""transpose the last 2 dimension of a tensor"""
|
|
x_new_dims = list(range(len(x.shape)))
|
|
x_new_dims[-1], x_new_dims[-2] = x_new_dims[-2], x_new_dims[-1]
|
|
x = paddle.transpose(x, x_new_dims)
|
|
return x
|
|
|
|
|
|
def get_inandout(A_shape, b_shape, trans="N", dtype="float64"):
|
|
paddle.disable_static(base.CPUPlace())
|
|
np.random.seed(2025)
|
|
A = np.random.random(A_shape).astype(dtype)
|
|
b = np.random.random(b_shape).astype(dtype)
|
|
x_grad_np = np.random.random(b_shape).astype(dtype)
|
|
if 'complex' in dtype:
|
|
A += 1j * np.random.random(A_shape).astype(dtype)
|
|
b += 1j * np.random.random(b_shape).astype(dtype)
|
|
x_grad_np += 1j * np.random.random(b_shape).astype(dtype)
|
|
x_grad = paddle.to_tensor(x_grad_np)
|
|
paddle_A = paddle.to_tensor(A)
|
|
lu, pivots = paddle.linalg.lu(paddle_A)
|
|
if trans == "N": # Ax = b
|
|
out = np.linalg.solve(A, b)
|
|
|
|
temp_A = np.swapaxes(A, -2, -1)
|
|
b_grad = np.linalg.solve(temp_A, x_grad)
|
|
|
|
_, L, U = paddle.linalg.lu_unpack(lu, pivots, True, False)
|
|
U_mH = _transpose_last_2dim(paddle.conj(U))
|
|
gR = paddle.linalg.triangular_solve(
|
|
U_mH,
|
|
paddle.mm(
|
|
-x_grad,
|
|
_transpose_last_2dim(paddle.conj(paddle.to_tensor(out))),
|
|
),
|
|
False,
|
|
False,
|
|
False,
|
|
)
|
|
gL = paddle.linalg.triangular_solve(
|
|
_transpose_last_2dim(paddle.conj(L)),
|
|
paddle.mm(gR, U_mH),
|
|
True,
|
|
False,
|
|
True,
|
|
)
|
|
lu_grad = (paddle.tril(gL, -1) + paddle.triu(gR, 0)).numpy()
|
|
elif trans == "T": # A^Tx = b
|
|
temp_A = np.swapaxes(A, -2, -1)
|
|
out = np.linalg.solve(temp_A, b)
|
|
|
|
b_grad = np.linalg.solve(A, x_grad)
|
|
|
|
P, L, U = paddle.linalg.lu_unpack(lu, pivots, True, True)
|
|
gR = paddle.mm(-_transpose_last_2dim(P), paddle.to_tensor(out))
|
|
gR = paddle.mm(gR, _transpose_last_2dim(paddle.conj(x_grad)))
|
|
gR = paddle.mm(gR, P)
|
|
L_mH = _transpose_last_2dim(paddle.conj(L))
|
|
gR = paddle.linalg.triangular_solve(L_mH, gR, True, True, True)
|
|
gU = paddle.linalg.triangular_solve(
|
|
_transpose_last_2dim(paddle.conj(U)),
|
|
paddle.mm(L_mH, gR),
|
|
False,
|
|
True,
|
|
False,
|
|
)
|
|
lu_grad = (paddle.tril(gR, -1) + paddle.triu(gU, 0)).numpy()
|
|
|
|
lu = lu.numpy()
|
|
pivots = pivots.numpy()
|
|
x_grad = x_grad.numpy()
|
|
paddle.enable_static()
|
|
return lu, pivots, b, out, x_grad, b_grad, lu_grad
|
|
|
|
|
|
class TestLuSolveOp(OpTest):
|
|
def setUp(self):
|
|
self.python_api = paddle.linalg.lu_solve
|
|
self.op_type = "lu_solve"
|
|
self.init_value()
|
|
(
|
|
self.LU,
|
|
self.pivots,
|
|
self.b,
|
|
self.out,
|
|
self.x_grad,
|
|
self.b_grad,
|
|
self.lu_grad,
|
|
) = get_inandout(self.A_shape, self.b_shape, self.trans, self.dtype)
|
|
self.inputs = {
|
|
'b': self.b,
|
|
'lu': self.LU,
|
|
'pivots': self.pivots,
|
|
}
|
|
self.attrs = {'trans': self.trans}
|
|
self.outputs = {'out': self.out}
|
|
|
|
def init_value(self):
|
|
self.A_shape = [2, 10, 10]
|
|
self.b_shape = [2, 10, 5]
|
|
self.trans = "N"
|
|
self.dtype = "float64"
|
|
|
|
def test_check_output(self):
|
|
paddle.enable_static()
|
|
self.check_output(check_pir=True)
|
|
paddle.disable_static()
|
|
|
|
def test_check_grad(self):
|
|
paddle.enable_static()
|
|
self.check_grad(
|
|
['b', 'lu'],
|
|
'out',
|
|
no_grad_set=['pivots'],
|
|
user_defined_grads=[self.b_grad, self.lu_grad],
|
|
user_defined_grad_outputs=[self.x_grad],
|
|
check_pir=True,
|
|
)
|
|
paddle.disable_static()
|
|
|
|
|
|
class TestLuSolveOp1(TestLuSolveOp):
|
|
def init_value(self):
|
|
self.A_shape = [2, 10, 10]
|
|
self.b_shape = [2, 10, 5]
|
|
self.trans = "T"
|
|
self.dtype = "float64"
|
|
|
|
|
|
class TestLuSolveOp2(TestLuSolveOp):
|
|
def init_value(self):
|
|
self.A_shape = [2, 2, 10, 10]
|
|
self.b_shape = [2, 2, 10, 5]
|
|
self.trans = "T"
|
|
self.dtype = "float64"
|
|
|
|
|
|
class TestLuSolveOp3(TestLuSolveOp):
|
|
def init_value(self):
|
|
self.A_shape = [2, 2, 10, 10]
|
|
self.b_shape = [2, 2, 10, 5]
|
|
self.trans = "N"
|
|
self.dtype = "float64"
|
|
|
|
|
|
class TestLuSolveOp4(TestLuSolveOp):
|
|
def init_value(self):
|
|
self.A_shape = [10, 10]
|
|
self.b_shape = [10, 10]
|
|
self.trans = "T"
|
|
self.dtype = "float64"
|
|
|
|
|
|
class TestLuSolveOp5(TestLuSolveOp):
|
|
def init_value(self):
|
|
self.A_shape = [10, 10]
|
|
self.b_shape = [10, 10]
|
|
self.trans = "N"
|
|
self.dtype = "float64"
|
|
|
|
|
|
# complex64
|
|
@unittest.skipIf(
|
|
base.core.is_compiled_with_rocm(), "Skip when compiled by ROCM."
|
|
)
|
|
class TestLuSolveOp6(TestLuSolveOp):
|
|
def init_value(self):
|
|
self.A_shape = [10, 10]
|
|
self.b_shape = [10, 10]
|
|
self.trans = "T"
|
|
self.dtype = "complex64"
|
|
|
|
|
|
# complex128
|
|
@unittest.skipIf(
|
|
base.core.is_compiled_with_rocm(), "Skip when compiled by ROCM."
|
|
)
|
|
class TestLuSolveOp7(TestLuSolveOp):
|
|
def init_value(self):
|
|
self.A_shape = [10, 10]
|
|
self.b_shape = [10, 10]
|
|
self.trans = "T"
|
|
self.dtype = "complex128"
|
|
|
|
|
|
class TestLuSolveOpAPI(unittest.TestCase):
|
|
def setUp(self):
|
|
self.init_value()
|
|
(
|
|
self.LU,
|
|
self.pivots,
|
|
self.b,
|
|
self.out,
|
|
_,
|
|
_,
|
|
_,
|
|
) = get_inandout(self.A_shape, self.b_shape, self.trans, self.dtype)
|
|
self.place = get_places()
|
|
|
|
def init_value(self):
|
|
# Ax = b
|
|
self.A_shape = [10, 10]
|
|
self.b_shape = [10, 5]
|
|
self.trans = "N"
|
|
self.dtype = "float64"
|
|
self.rtol = 1e-05
|
|
|
|
def test_dygraph(self):
|
|
def run(place):
|
|
paddle.disable_static(place)
|
|
lu = paddle.to_tensor(self.LU)
|
|
pivots = paddle.to_tensor(self.pivots)
|
|
b = paddle.to_tensor(self.b)
|
|
lu_solve_x = paddle.linalg.lu_solve(b, lu, pivots, self.trans)
|
|
np.testing.assert_allclose(
|
|
lu_solve_x.numpy(), self.out, rtol=self.rtol
|
|
)
|
|
paddle.enable_static()
|
|
|
|
for place in self.place:
|
|
run(place)
|
|
|
|
def test_static(self):
|
|
def run(place):
|
|
paddle.enable_static()
|
|
with paddle.static.program_guard(
|
|
paddle.static.Program(), paddle.static.Program()
|
|
):
|
|
b = paddle.static.data(
|
|
name='B', shape=self.b.shape, dtype=self.b.dtype
|
|
)
|
|
lu = paddle.static.data(
|
|
name='Lu', shape=self.LU.shape, dtype=self.LU.dtype
|
|
)
|
|
pivots = paddle.static.data(
|
|
name='Pivots',
|
|
shape=self.pivots.shape,
|
|
dtype=self.pivots.dtype,
|
|
)
|
|
lu_solve_x = paddle.linalg.lu_solve(b, lu, pivots, self.trans)
|
|
exe = base.Executor(place)
|
|
fetches = exe.run(
|
|
feed={
|
|
'B': self.b,
|
|
'Lu': self.LU,
|
|
'Pivots': self.pivots,
|
|
},
|
|
fetch_list=[lu_solve_x],
|
|
)
|
|
np.testing.assert_allclose(fetches[0], self.out, rtol=self.rtol)
|
|
paddle.disable_static()
|
|
|
|
for place in self.place:
|
|
run(place)
|
|
|
|
|
|
class TestLuSolveOpAPI2(TestLuSolveOpAPI):
|
|
def init_value(self):
|
|
# Ax = b
|
|
self.A_shape = [1, 10, 10]
|
|
self.b_shape = [2, 10, 5]
|
|
self.trans = "N"
|
|
self.dtype = "float64"
|
|
self.rtol = 1e-05
|
|
|
|
|
|
class TestLuSolveOpAPI3(TestLuSolveOpAPI):
|
|
def init_value(self):
|
|
# A^Tx = b
|
|
self.A_shape = [1, 10, 10]
|
|
self.b_shape = [2, 10, 5]
|
|
self.trans = "T"
|
|
self.dtype = "float64"
|
|
self.rtol = 1e-05
|
|
|
|
|
|
class TestLuSolveOpAPI4(TestLuSolveOpAPI):
|
|
def init_value(self):
|
|
# Ax = b
|
|
self.A_shape = [1, 10, 10]
|
|
self.b_shape = [2, 10, 5]
|
|
self.trans = "N"
|
|
self.dtype = "float32"
|
|
self.rtol = 0.001
|
|
|
|
|
|
class TestLuSolveOpAPI5(TestLuSolveOpAPI):
|
|
def init_value(self):
|
|
# A^Tx = b
|
|
self.A_shape = [1, 10, 10]
|
|
self.b_shape = [2, 10, 5]
|
|
self.trans = "T"
|
|
self.dtype = "float32"
|
|
self.rtol = 0.001
|
|
|
|
|
|
class TestLuSolveOpAPI6(TestLuSolveOpAPI):
|
|
def init_value(self):
|
|
# Ax = b
|
|
self.A_shape = [10, 10]
|
|
self.b_shape = [10, 5]
|
|
self.trans = "N"
|
|
self.dtype = "float32"
|
|
self.rtol = 0.001
|
|
|
|
|
|
class TestLuSolveOpAPI7(TestLuSolveOpAPI):
|
|
def init_value(self):
|
|
# A^Tx = b
|
|
self.A_shape = [10, 10]
|
|
self.b_shape = [10, 5]
|
|
self.trans = "T"
|
|
self.dtype = "float32"
|
|
self.rtol = 0.001
|
|
|
|
|
|
class TestLuSolveOpAPI8(TestLuSolveOpAPI):
|
|
def init_value(self):
|
|
# A^Tx = b
|
|
self.A_shape = [10, 10]
|
|
self.b_shape = [10, 5]
|
|
self.trans = "T"
|
|
self.dtype = "float64"
|
|
self.rtol = 1e-05
|
|
|
|
|
|
@unittest.skipIf(
|
|
base.core.is_compiled_with_rocm(), "Skip when compiled by ROCM."
|
|
)
|
|
class TestLuSolveOpAPI9(TestLuSolveOpAPI):
|
|
def init_value(self):
|
|
# Ax = b
|
|
self.A_shape = [10, 10]
|
|
self.b_shape = [10, 5]
|
|
self.trans = "N"
|
|
self.dtype = "complex64"
|
|
self.rtol = 0.001
|
|
|
|
|
|
@unittest.skipIf(
|
|
base.core.is_compiled_with_rocm(), "Skip when compiled by ROCM."
|
|
)
|
|
class TestLuSolveOpAPI10(TestLuSolveOpAPI):
|
|
def init_value(self):
|
|
# Ax = b
|
|
self.A_shape = [10, 10]
|
|
self.b_shape = [10, 5]
|
|
self.trans = "N"
|
|
self.dtype = "complex128"
|
|
self.rtol = 1e-05
|
|
|
|
|
|
class TestLSolveError(unittest.TestCase):
|
|
def test_errors(self):
|
|
with paddle.base.dygraph.guard():
|
|
# The size of b should gather than 2.
|
|
def test_b_size():
|
|
b = paddle.randn([3])
|
|
lu = paddle.randn([3, 3])
|
|
pivots = paddle.randn([3])
|
|
paddle.linalg.lu_solve(b, lu, pivots)
|
|
|
|
self.assertRaises(ValueError, test_b_size)
|
|
|
|
# The size of lu should gather than 2.
|
|
def test_lu_size():
|
|
b = paddle.randn([3, 1])
|
|
lu = paddle.randn([3])
|
|
pivots = paddle.randn([3])
|
|
paddle.linalg.lu_solve(b, lu, pivots)
|
|
|
|
self.assertRaises(ValueError, test_lu_size)
|
|
|
|
# The size of pivots should gather than 1.
|
|
def test_pivots_size():
|
|
b = paddle.randn([3, 1])
|
|
lu = paddle.randn([3, 3])
|
|
pivots = paddle.randn([])
|
|
paddle.linalg.lu_solve(b, lu, pivots)
|
|
|
|
self.assertRaises(ValueError, test_pivots_size)
|
|
|
|
# b.shape[-2] should equal to lu.shape[-2].
|
|
def test_b_lu_shape():
|
|
b = paddle.randn([1, 3])
|
|
lu = paddle.randn([3, 3])
|
|
pivots = paddle.randn([3])
|
|
paddle.linalg.lu_solve(b, lu, pivots)
|
|
|
|
self.assertRaises(ValueError, test_b_lu_shape)
|
|
|
|
# lu.shape[-1] should equal to pivots.shape[-1].
|
|
def test_b_pivots_shape():
|
|
b = paddle.randn([3, 1])
|
|
lu = paddle.randn([3, 3])
|
|
pivots = paddle.randn([2])
|
|
paddle.linalg.lu_solve(b, lu, pivots)
|
|
|
|
self.assertRaises(ValueError, test_b_pivots_shape)
|
|
|
|
# lu.shape[-2] should equal to lu.shape[-1].
|
|
def test_lu_shape():
|
|
b = paddle.randn([3, 1])
|
|
lu = paddle.randn([3, 2])
|
|
pivots = paddle.randn([3])
|
|
paddle.linalg.lu_solve(b, lu, pivots)
|
|
|
|
self.assertRaises(ValueError, test_lu_shape)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|