jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_lu_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value

        A = random_fullrank_matrix_distinct_singular_value(2, 2, dtype=dtype, device='cpu')
        b = torch.randn(2, 2, 2, dtype=dtype, device='cpu')
        x_exp = torch.as_tensor(solve(A.permute(0, 2, 1).numpy(), b.permute(2, 1, 0).numpy())).to(device)
        A = A.to(device).permute(0, 2, 1)
        b = b.to(device).permute(2, 1, 0)
        assert not A.is_contiguous() and not b.is_contiguous(), "contiguous inputs"
        LU_data, LU_pivots = torch.lu(A)
        x = torch.lu_solve(b, LU_data, LU_pivots)
        self.assertEqual(x, x_exp)
//  [python]
def test_lu_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value

        A = random_fullrank_matrix_distinct_singular_value(2, 2, dtype=dtype, device='cpu')
        b = torch.randn(2, 2, 2, dtype=dtype, device='cpu')
        x_exp = torch.as_tensor(solve(A.permute(0, 2, 1).numpy(), b.permute(2, 1, 0).numpy())).to(device)
        A = A.to(device).permute(0, 2, 1)
        b = b.to(device).permute(2, 1, 0)
        assert not A.is_contiguous() and not b.is_contiguous(), "contiguous inputs"
        LU_data, LU_pivots = torch.linalg.lu_factor(A)
        x = torch.lu_solve(b, LU_data, LU_pivots)
        self.assertEqual(x, x_exp)
//  [python]
def weight_inverse_and_logabsdet(self):
        """
        Cost:
            inverse = O(D^3)
            logabsdet = O(D)
        where:
            D = num of features
        """
        # If both weight inverse and logabsdet are needed, it's cheaper to compute both together.
        identity = torch.eye(self.features, self.features)
        # LU-decompose the weights and solve for the outputs.
        lu, lu_pivots = torch.lu(self._weight)
        weight_inv = torch.lu_solve(identity, lu, lu_pivots)
        logabsdet = torch.sum(torch.log(torch.abs(torch.diag(lu))))
        return weight_inv, logabsdet
//  [python]
def weight_inverse_and_logabsdet(self):
        """
        Cost:
            inverse = O(D^3)
            logabsdet = O(D)
        where:
            D = num of features
        """
        # If both weight inverse and logabsdet are needed, it's cheaper to compute both together.
        identity = torch.eye(self.features, self.features)
        # LU-decompose the weights and solve for the outputs.
        lu, pivots = torch.linalg.lu_factor(self._weight)
        weight_inv = torch.linalg.lu_solve(identity, lu, pivots)
        logabsdet = torch.sum(torch.log(torch.abs(torch.diag(lu))))
        return weight_inv, logabsdet
//  [python]
def test_lu_unpack_check_input(self, device, dtype):
        x = torch.rand(5, 5, 5, device=device, dtype=dtype)
        lu_data, lu_pivots = torch.lu(x, pivot=True)

        with self.assertRaisesRegex(RuntimeError, "torch.int32 dtype"):
            torch.lu_unpack(lu_data, lu_pivots.long())
        with self.assertRaisesRegex(RuntimeError, "contiguous tensor"):
            torch.lu_unpack(lu_data, lu_pivots.transpose(-1, -2))

        # check that onces flags are unset, Nones are returned
        p, l, u = torch.lu_unpack(lu_data, lu_pivots, unpack_data=False)
        self.assertTrue((l == u) and l is None)
        p, l, u = torch.lu_unpack(lu_data, lu_pivots, unpack_pivots=False)
        self.assertTrue(p is None)
        p, l, u = torch.lu_unpack(lu_data, lu_pivots, unpack_data=False, unpack_pivots=False)
        self.assertTrue((p == l == u) and p is None)
//  [python]
def test_lu_unpack_check_input(self, device, dtype):
        x = torch.rand(5, 5, 5, device=device, dtype=dtype)
        lu_data, lu_pivots = torch.linalg.lu_factor(x, pivot=True)

        with self.assertRaisesRegex(RuntimeError, "torch.int32 dtype"):
            torch.lu_unpack(lu_data, lu_pivots.long())
        with self.assertRaisesRegex(RuntimeError, "contiguous tensor"):
            torch.lu_unpack(lu_data, lu_pivots.transpose(-1, -2))

        # check that onces flags are unset, Nones are returned
        p, l, u = torch.lu_unpack(lu_data, lu_pivots, unpack_data=False)
        self.assertTrue((l == u) and l is None)
        p, l, u = torch.lu_unpack(lu_data, lu_pivots, unpack_pivots=False)
        self.assertTrue(p is None)
        p, l, u = torch.lu_unpack(lu_data, lu_pivots, unpack_data=False, unpack_pivots=False)
        self.assertTrue((p == l == u) and p is None)
//  [python]
def test_lu_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value

        def run_test(A_dims, b_dims, pivot=True):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = random_fullrank_matrix_distinct_singular_value(A_matrix_size, *A_batch_dims, dtype=dtype)
            b = torch.randn(*b_dims, dtype=dtype)
            x_exp = torch.as_tensor(solve(A.numpy(), b.numpy())).to(dtype=dtype, device=device)
            A, b = A.to(device), b.to(device)
            LU_data, LU_pivots = torch.lu(A, pivot=pivot)
            x = torch.lu_solve(b, LU_data, LU_pivots)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6))  # no broadcasting
        run_test((2, 1, 3, 4, 4), (4, 6))  # broadcasting b
        run_test((4, 4), (2, 1, 3, 4, 2))  # broadcasting A
        run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5))
//  [python]
def test_lu_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value

        def run_test(A_dims, b_dims, pivot=True):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = random_fullrank_matrix_distinct_singular_value(A_matrix_size, *A_batch_dims, dtype=dtype)
            b = torch.randn(*b_dims, dtype=dtype)
            x_exp = torch.as_tensor(solve(A.numpy(), b.numpy())).to(dtype=dtype, device=device)
            A, b = A.to(device), b.to(device)
            LU_data, LU_pivots = torch.linalg.lu_factor(A, pivot=pivot)
            x = torch.lu_solve(b, LU_data, LU_pivots)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6))  # no broadcasting
        run_test((2, 1, 3, 4, 4), (4, 6))  # broadcasting b
        run_test((4, 4), (2, 1, 3, 4, 2))  # broadcasting A
        run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5))
//  [python]
def test_lu_solve_batched(self, device, dtype):
        def sub_test(pivot):
            def lu_solve_batch_test_helper(A_dims, b_dims, pivot):
                b, A, LU_data, LU_pivots = self.lu_solve_test_helper(A_dims, b_dims, pivot, device, dtype)
                x_exp_list = []
                for i in range(b_dims[0]):
                    x_exp_list.append(torch.lu_solve(b[i], LU_data[i], LU_pivots[i]))
                x_exp = torch.stack(x_exp_list)  # Stacked output
                x_act = torch.lu_solve(b, LU_data, LU_pivots)  # Actual output
                self.assertEqual(x_exp, x_act)  # Equality check
                Ax = torch.matmul(A, x_act)
                self.assertEqual(b, Ax)

            for batchsize in [1, 3, 4]:
                lu_solve_batch_test_helper((5, batchsize), (batchsize, 5, 10), pivot)

        # Tests tensors with 0 elements
        b = torch.randn(3, 0, 3, dtype=dtype, device=device)
        A = torch.randn(3, 0, 0, dtype=dtype, device=device)
        LU_data, LU_pivots = torch.lu(A)
        self.assertEqual(torch.empty_like(b), b.lu_solve(LU_data, LU_pivots))

        sub_test(True)
        if self.device_type == 'cuda':
            sub_test(False)
//  [python]
def test_lu_solve_batched(self, device, dtype):
        def sub_test(pivot):
            def lu_solve_batch_test_helper(A_dims, b_dims, pivot):
                b, A, LU_data, LU_pivots = self.lu_solve_test_helper(A_dims, b_dims, pivot, device, dtype)
                x_exp_list = []
                for i in range(b_dims[0]):
                    x_exp_list.append(torch.linalg.solve(b[i], torch.linalg.lu_factor(A[i], pivot=pivot)))
                x_exp = torch.stack(x_exp_list)  # Stacked output
                x_act = torch.linalg.solve(b, torch.linalg.lu_factor(A, pivot=pivot))  # Actual output
                self.assertEqual(x_exp, x_act)  # Equality check
                Ax = torch.matmul(A, x_act)
                self.assertEqual(b, Ax)

            for batchsize in [1, 3, 4]:
                lu_solve_batch_test_helper((5, batchsize), (batchsize, 5, 10), pivot)

        # Tests tensors with 0 elements
        b = torch.randn(3, 0, 3, dtype=dtype, device=device)
        A = torch.randn(3, 0, 0, dtype=dtype, device=device)
        LU_data, LU_pivots = torch.linalg.lu_factor(A)
        self.assertEqual(torch.empty_like(b), torch.linalg.solve(b, LU_data))

        sub_test(True)
        if self.device_type == 'cuda':
            sub_test(False)
//  [python]
def test_lu_solve_out_errors_and_warnings(self, device, dtype):
        # dtypes should be safely castable
        a = torch.eye(2, dtype=dtype, device=device)
        LU_data, LU_pivots = torch.lu(a, pivot=True)
        b = torch.randn(2, 1, dtype=dtype, device=device)
        out = torch.empty(0, dtype=torch.int, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got result with dtype Int"):
            torch.lu_solve(b, LU_data, LU_pivots, out=out)

        # device should match
        if torch.cuda.is_available():
            wrong_device = 'cpu' if self.device_type != 'cpu' else 'cuda'
            out = torch.empty(0, dtype=dtype, device=wrong_device)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.lu_solve(b, LU_data, LU_pivots, out=out)

        # if out tensor with wrong shape is passed a warning is given
        with warnings.catch_warnings(record=True) as w:
            out = torch.empty(1, dtype=dtype, device=device)
            # Trigger warning
            torch.lu_solve(b, LU_data, LU_pivots, out=out)
            # Check warning occurs
            self.assertEqual(len(w), 1)
            self.assertTrue("An output with one or more elements was resized" in str(w[-1].message))
//  [python]
def test_lu_solve_out_errors_and_warnings(self, device, dtype):
        # dtypes should be safely castable
        a = torch.eye(2, dtype=dtype, device=device)
        LU_data, LU_pivots = torch.linalg.lu_factor(a, pivot=True)
        b = torch.randn(2, 1, dtype=dtype, device=device)
        out = torch.empty(0, dtype=torch.int, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got result with dtype Int"):
            torch.lu_solve(b, LU_data, LU_pivots, out=out)

        # device should match
        if torch.cuda.is_available():
            wrong_device = 'cpu' if self.device_type != 'cpu' else 'cuda'
            out = torch.empty(0, dtype=dtype, device=wrong_device)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.lu_solve(b, LU_data, LU_pivots, out=out)

        # if out tensor with wrong shape is passed a warning is given
        with warnings.catch_warnings(record=True) as w:
            out = torch.empty(1, dtype=dtype, device=device)
            # Trigger warning
            torch.lu_solve(b, LU_data, LU_pivots, out=out)
            # Check warning occurs
            self.assertEqual(len(w), 1)
            self.assertTrue("An output with one or more elements was resized" in str(w[-1].message))
//  [python]
def test_lu_solve_batched(self, device, dtype):
        def sub_test(pivot):
            def lu_solve_batch_test_helper(A_dims, b_dims, pivot):
                b, A, LU_data, LU_pivots = self.lu_solve_test_helper(A_dims, b_dims, pivot, device, dtype)
                x_exp_list = []
                for i in range(b_dims[0]):
                    x_exp_list.append(torch.lu_solve(b[i], LU_data[i], LU_pivots[i]))
                x_exp = torch.stack(x_exp_list)  # Stacked output
                x_act = torch.lu_solve(b, LU_data, LU_pivots)  # Actual output
                self.assertEqual(x_exp, x_act)  # Equality check
                Ax = np.matmul(A.cpu(), x_act.cpu())
                self.assertEqual(b, Ax)

            for batchsize in [1, 3, 4]:
                lu_solve_batch_test_helper((5, batchsize), (batchsize, 5, 10), pivot)

        # Tests tensors with 0 elements
        b = torch.randn(3, 0, 3, dtype=dtype, device=device)
        A = torch.randn(3, 0, 0, dtype=dtype, device=device)
        LU_data, LU_pivots = torch.lu(A)
        self.assertEqual(torch.empty_like(b), b.lu_solve(LU_data, LU_pivots))

        sub_test(True)
        if self.device_type == 'cuda':
            sub_test(False)
//  [python]
def test_lu_solve_batched(self, device, dtype):
        def sub_test(pivot):
            def lu_solve_batch_test_helper(A_dims, b_dims, pivot):
                b, A, LU_data, LU_pivots = self.lu_solve_test_helper(A_dims, b_dims, pivot, device, dtype)
                x_exp_list = []
                for i in range(b_dims[0]):
                    x_exp_list.append(torch.linalg.solve(b[i], torch.linalg.lu_factor(A[i], pivot=pivot)))
                x_exp = torch.stack(x_exp_list)  # Stacked output
                x_act = torch.linalg.solve(b, torch.linalg.lu_factor(A, pivot=pivot))  # Actual output
                self.assertEqual(x_exp, x_act)  # Equality check
                Ax = np.matmul(A.cpu(), x_act.cpu())
                self.assertEqual(b, Ax)

            for batchsize in [1, 3, 4]:
                lu_solve_batch_test_helper((5, batchsize), (batchsize, 5, 10), pivot)

        # Tests tensors with 0 elements
        b = torch.randn(3, 0, 3, dtype=dtype, device=device)
        A = torch.randn(3, 0, 0, dtype=dtype, device=device)
        LU_data, LU_pivots = torch.linalg.lu_factor(A)
        self.assertEqual(torch.empty_like(b), torch.linalg.solve(b, LU_data))

        sub_test(True)
        if self.device_type == 'cuda':
            sub_test(False)
//  [python]
def test_dcem():
    n_batch = 2
    n_sample = 100
    N = 2

    torch.manual_seed(0)
    Q = torch.eye(N).unsqueeze(0).repeat(n_batch, 1, 1)
    p = 0.1*torch.randn(n_batch, N)

    Q_sample = Q.unsqueeze(1).repeat(1, n_sample, 1, 1).view(n_batch*n_sample, N, N)
    p_sample = p.unsqueeze(1).repeat(1, n_sample, 1).view(n_batch*n_sample, N)

    def f(X):
        assert X.size() == (n_batch, n_sample, N)
        X = X.view(n_batch*n_sample, N)
        obj = 0.5*(bmv(Q_sample, X)*X).sum(dim=1) + (p_sample*X).sum(dim=1)
        obj = obj.view(n_batch, n_sample)
        return obj

    def iter_cb(i, X, fX, I, X_I, mu, sigma):
        print(fX.mean(dim=1))

    xhat = dcem(
        f, nx = N, n_batch = n_batch, n_sample = n_sample,
        n_elite = 50, n_iter = 40,
        temp = 1.,
        normalize = True,
        # temp = np.infty,
        init_sigma=1.,
        iter_cb = iter_cb,
        # lb = -5., ub = 5.,
    )

    Q_LU = torch.lu(Q)
    xstar = -torch.lu_solve(p, *Q_LU)
    assert (xhat-xstar).abs().max() < 1e-4
//  [python]
def test_dcem():
    n_batch = 2
    n_sample = 100
    N = 2

    torch.manual_seed(0)
    Q = torch.eye(N).unsqueeze(0).repeat(n_batch, 1, 1)
    p = 0.1*torch.randn(n_batch, N)

    Q_sample = Q.unsqueeze(1).repeat(1, n_sample, 1, 1).view(n_batch*n_sample, N, N)
    p_sample = p.unsqueeze(1).repeat(1, n_sample, 1).view(n_batch*n_sample, N)

    def f(X):
        assert X.size() == (n_batch, n_sample, N)
        X = X.view(n_batch*n_sample, N)
        obj = 0.5*(bmv(Q_sample, X)*X).sum(dim=1) + (p_sample*X).sum(dim=1)
        obj = obj.view(n_batch, n_sample)
        return obj

    def iter_cb(i, X, fX, I, X_I, mu, sigma):
        print(fX.mean(dim=1))

    xhat = dcem(
        f, nx = N, n_batch = n_batch, n_sample = n_sample,
        n_elite = 50, n_iter = 40,
        temp = 1.,
        normalize = True,
        # temp = np.infty,
        init_sigma=1.,
        iter_cb = iter_cb,
        # lb = -5., ub = 5.,
    )

    Q_LU = torch.linalg.lu_factor(Q)
    xstar = -torch.linalg.lu_solve(p, *Q_LU)
    assert (xhat-xstar).abs().max() < 1e-4
//  [python]
def _cayley(A):
    I = torch.eye(A.shape[-1], device=A.device)
    LU  = torch.lu(I+A, pivot=True)
    return torch.lu_solve(I-A,*LU)
//  [python]
def _cayley(A):
    I = torch.eye(A.shape[-1], device=A.device)
    LU  = torch.linalg.lu_factor(I+A, pivot=True)
    return torch.lu_solve(I-A,*LU)
//  [python]
def __init__(self, in_out_channels):
        super(InvertibleConv1x1, self).__init__()

        W = torch.zeros((in_out_channels, in_out_channels), dtype=torch.float32)
        nn.init.orthogonal_(W)
        LU, pivots = torch.lu(W)

        P, L, U = torch.lu_unpack(LU, pivots)
        self.P = nn.Parameter(P, requires_grad=False)
        self.L = nn.Parameter(L, requires_grad=True)
        self.U = nn.Parameter(U, requires_grad=True)
        self.I = nn.Parameter(torch.eye(in_out_channels), requires_grad=False)
        self.pivots = nn.Parameter(pivots, requires_grad=False)

        L_mask = np.tril(np.ones((in_out_channels, in_out_channels), dtype='float32'), k=-1)
        U_mask = L_mask.T.copy()
        self.L_mask = nn.Parameter(torch.from_numpy(L_mask), requires_grad=False)
        self.U_mask = nn.Parameter(torch.from_numpy(U_mask), requires_grad=False)

        s = torch.diag(U)
        sign_s = torch.sign(s)
        log_s = torch.log(torch.abs(s))
        self.log_s = nn.Parameter(log_s, requires_grad=True)
        self.sign_s = nn.Parameter(sign_s, requires_grad=False)
//  [python]
def __init__(self, in_out_channels):
        super(InvertibleConv1x1, self).__init__()

        W = torch.zeros((in_out_channels, in_out_channels), dtype=torch.float32)
        nn.init.orthogonal_(W)
        LU, pivots = torch.linalg.lu_factor(W)

        P, L, U = torch.lu_unpack(LU, pivots)
        self.P = nn.Parameter(P, requires_grad=False)
        self.L = nn.Parameter(L, requires_grad=True)
        self.U = nn.Parameter(U, requires_grad=True)
        self.I = nn.Parameter(torch.eye(in_out_channels), requires_grad=False)
        self.pivots = nn.Parameter(pivots, requires_grad=False)

        L_mask = np.tril(np.ones((in_out_channels, in_out_channels), dtype='float32'), k=-1)
        U_mask = L_mask.T.copy()
        self.L_mask = nn.Parameter(torch.from_numpy(L_mask), requires_grad=False)
        self.U_mask = nn.Parameter(torch.from_numpy(U_mask), requires_grad=False)

        s = torch.diag(U)
        sign_s = torch.sign(s)
        log_s = torch.log(torch.abs(s))
        self.log_s = nn.Parameter(log_s, requires_grad=True)
        self.sign_s = nn.Parameter(sign_s, requires_grad=False)
//  [python]
def _cayley_inverse(Q):
    I = torch.eye(Q.shape[-1], device=Q.device)
    rec_LU = torch.lu(I+Q, pivot=True)
    return torch.lu_solve(I-Q,*rec_LU)
//  [python]
def _cayley_inverse(Q):
    I = torch.eye(Q.shape[-1], device=Q.device)
    rec_LU = torch.linalg.lu_factor(I+Q, pivot=True)
    return torch.linalg.lu_solve(I-Q,*rec_LU)
//  [python]
def _cayley_frechet(A,H,Q=None):
    I = torch.eye(A.shape[-1], device=A.device)
    if Q is None:
        Q = _cayley(A)
    _LU = torch.lu(I+A, pivot=True)
    p = torch.lu_solve(Q, *_LU)
    _LU = torch.lu(I-A, pivot=True)
    q = torch.lu_solve(H, *_LU)
    return 2.* q @ p
//  [python]
def _cayley_frechet(A,H,Q=None):
    I = torch.eye(A.shape[-1], device=A.device)
    if Q is None:
        Q = _cayley(A)
    _LU = torch.linalg.lu_factor(I+A, pivot=True)
    p = torch.linalg.lu_solve(Q, *_LU)
    _LU = torch.linalg.lu_factor(I-A, pivot=True)
    q = torch.linalg.lu_solve(H, *_LU)
    return 2.* q @ p
//  [python]
def xpbatch_lu_factor(A):
    assert len(A.shape) == 3, "Actual" + str(A.shape)
    assert A.shape[1] == A.shape[2], "Actual" + str(A.shape)
    xp = get_array_module(A)
    A = copy.deepcopy(A)
    '''
    if cupy_available and xp == cupy:
        Ps = xp.empty((A.shape[0], A.shape[1]), dtype=np.int32)
        for i in range(len(A)):
            A[i], Ps[i] = cupyx.scipy.linalg.u_factor(A[i], overwrite_a=True)
        return A, Ps
    else:
        Ps = []
        for i in range(len(A)):
            A[i], piv = scipy.linalg.lu_factor(A[i], overwrite_a=True)
            Ps.append(piv)
        return A, Ps
    '''
    # use PyTorch here, because scipy does not offer batch lu_factorization
    A_LU, pivots = torch.lu(torch.tensor(A))
    return A_LU.cpu().numpy(), pivots.cpu().numpy()
//  [python]
def xpbatch_lu_factor(A):
    assert len(A.shape) == 3, "Actual" + str(A.shape)
    assert A.shape[1] == A.shape[2], "Actual" + str(A.shape)
    xp = get_array_module(A)
    A = copy.deepcopy(A)
    A_LU, pivots = torch.linalg.lu_factor(torch.tensor(A))
    return A_LU.cpu().numpy(), pivots.cpu().numpy()
//  [python]
def forward(ctx, A, b):
        A_LU, pivots = torch.lu(A)
        x = torch.lu_solve(b, A_LU, pivots)

        ctx.save_for_backward(A_LU, pivots, x)
        
        return x
//  [python]
def forward(ctx, A, b):
        A_LU, pivots = torch.linalg.lu_factor(A)
        x = torch.lu_solve(b, A_LU, pivots)

        ctx.save_for_backward(A_LU, pivots, x)
        
        return x
//  [python]
def lu_solve_test_helper(self, A_dims, b_dims, pivot, device, dtype):
        make_fullrank = make_fullrank_matrices_with_distinct_singular_values
        make_A = partial(make_fullrank, device=device, dtype=dtype)

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = make_A(*A_dims)
        LU_data, LU_pivots, info = torch.lu(A, get_infos=True, pivot=pivot)
        self.assertEqual(info, torch.zeros_like(info))
        return b, A, LU_data, LU_pivots
//  [python]
def lu_solve_test_helper(self, A_dims, b_dims, pivot, device, dtype):
        make_fullrank = make_fullrank_matrices_with_distinct_singular_values
        make_A = partial(make_fullrank, device=device, dtype=dtype)

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = make_A(*A_dims)
        LU_data, LU_pivots = torch.linalg.lu_factor(A, pivot=pivot)
        self.assertEqual(LU_pivots, torch.zeros_like(LU_pivots))
        return b, A, LU_data, LU_pivots
//  [python]
def safe_solve_with_mask(B: Tensor, A: Tensor) -> tuple[Tensor, Tensor, Tensor]:
	r"""Helper function, which avoids crashing because of singular matrix input
	and outputs the mask of valid solution
	"""
	if not torch_version_geq(1, 10):
		sol, lu = _torch_solve_cast(B, A)
		error_console.log("PyTorch version < 1.10, solve validness mask maybe not "
						  "correct", RuntimeWarning)
		return sol, lu, torch.ones(len(A), dtype=torch.bool, device=A.device)
	# Based on https://github.com/pytorch/pytorch/issues/31546#issuecomment-694135622
	if not isinstance(B, Tensor):
		raise AssertionError(f"B must be Tensor. Got: {type(B)}.")
	dtype = B.dtype
	if dtype not in (torch.float32, torch.float64):
		dtype = torch.float32
	A_LU, pivots, info = torch.lu(A.to(dtype), get_infos=True)
	valid_mask         = info == 0
	X = torch.lu_solve(B.to(dtype), A_LU, pivots)
	return X.to(B.dtype), A_LU.to(A.dtype), valid_mask
//  [python]
def safe_solve_with_mask(B: Tensor, A: Tensor) -> tuple[Tensor, Tensor, Tensor]:
	r"""Helper function, which avoids crashing because of singular matrix input
	and outputs the mask of valid solution
	"""
	if not torch_version_geq(1, 10):
		sol, lu = _torch_solve_cast(B, A)
		error_console.log("PyTorch version < 1.10, solve validness mask maybe not "
						  "correct", RuntimeWarning)
		return sol, lu, torch.ones(len(A), dtype=torch.bool, device=A.device)
	# Based on https://github.com/pytorch/pytorch/issues/31546#issuecomment-694135622
	if not isinstance(B, Tensor):
		raise AssertionError(f"B must be Tensor. Got: {type(B)}.")
	dtype = B.dtype
	if dtype not in (torch.float32, torch.float64):
		dtype = torch.float32
	A_LU, pivots = torch.linalg.lu_factor(A.to(dtype))
	valid_mask         = pivots[:,0] >= 0
	X = torch.linalg.lu_solve(B.to(dtype), A_LU, pivots)
	return X.to(B.dtype), A_LU.to(A.dtype), valid_mask
//  [python]
def safe_solve_with_mask(B: torch.Tensor, A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    r"""Helper function, which avoids crashing because of singular matrix input and outputs the
    mask of valid solution"""
    if not torch_version_geq(1, 10):
        sol, lu = _torch_solve_cast(B, A)
        warnings.warn('PyTorch version < 1.10, solve validness mask maybe not correct', RuntimeWarning)
        return sol, lu, torch.ones(len(A), dtype=torch.bool, device=A.device)
    # Based on https://github.com/pytorch/pytorch/issues/31546#issuecomment-694135622
    if not isinstance(B, torch.Tensor):
        raise AssertionError(f"B must be torch.Tensor. Got: {type(B)}.")
    dtype: torch.dtype = B.dtype
    if dtype not in (torch.float32, torch.float64):
        dtype = torch.float32
    A_LU, pivots, info = torch.lu(A.to(dtype), get_infos=True)
    valid_mask: torch.Tensor = info == 0
    X = torch.lu_solve(B.to(dtype), A_LU, pivots)
    return X.to(B.dtype), A_LU.to(A.dtype), valid_mask
//  [python]
def safe_solve_with_mask(B: torch.Tensor, A: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    r"""Helper function, which avoids crashing because of singular matrix input and outputs the
    mask of valid solution"""
    if not torch_version_geq(1, 10):
        sol, lu = _torch_solve_cast(B, A)
        warnings.warn('PyTorch version < 1.10, solve validness mask maybe not correct', RuntimeWarning)
        return sol, lu, torch.ones(len(A), dtype=torch.bool, device=A.device)
    # Based on https://github.com/pytorch/pytorch/issues/31546#issuecomment-694135622
    if not isinstance(B, torch.Tensor):
        raise AssertionError(f"B must be torch.Tensor. Got: {type(B)}.")
    dtype: torch.dtype = B.dtype
    if dtype not in (torch.float32, torch.float64):
        dtype = torch.float32
    A_LU, pivots = torch.linalg.lu_factor(A.to(dtype))
    valid_mask: torch.Tensor = (torch.abs(torch.diag(A_LU)) > 1e-10).all()
    X = torch.linalg.lu_solve(B.to(dtype), A_LU, pivots)
    return X.to(B.dtype), A_LU.to(A.dtype), valid_mask
//  [python]
def test_lu_solve_batched_broadcasting(self, device, dtype):
        make_fullrank = make_fullrank_matrices_with_distinct_singular_values
        make_A = partial(make_fullrank, device=device, dtype=dtype)

        def run_test(A_dims, b_dims, pivot=True):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = make_A(*A_batch_dims, A_matrix_size, A_matrix_size)
            b = make_tensor(b_dims, dtype=dtype, device=device)
            x_exp = np.linalg.solve(A.cpu(), b.cpu())
            LU_data, LU_pivots = torch.lu(A, pivot=pivot)
            x = torch.lu_solve(b, LU_data, LU_pivots)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6))  # no broadcasting
        run_test((2, 1, 3, 4, 4), (4, 6))  # broadcasting b
        run_test((4, 4), (2, 1, 3, 4, 2))  # broadcasting A
        run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5))
//  [python]
def test_lu_solve_batched_broadcasting(self, device, dtype):
        make_fullrank = make_fullrank_matrices_with_distinct_singular_values
        make_A = partial(make_fullrank, device=device, dtype=dtype)

        def run_test(A_dims, b_dims, pivot=True):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = make_A(*A_batch_dims, A_matrix_size, A_matrix_size)
            b = make_tensor(b_dims, dtype=dtype, device=device)
            x_exp = np.linalg.solve(A.cpu(), b.cpu())
            LU_data, LU_pivots = torch.linalg.lu_factor(A, pivot=pivot)
            x = torch.linalg.lu_solve(b, LU_data, LU_pivots)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6))  # no broadcasting
        run_test((2, 1, 3, 4, 4), (4, 6))  # broadcasting b
        run_test((4, 4), (2, 1, 3, 4, 2))  # broadcasting A
        run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5))
//  [python]
def test_lu_solve_batched(self, device, dtype):
        def sub_test(pivot):
            def lu_solve_batch_test_helper(A_dims, b_dims, pivot):
                b, A, LU_data, LU_pivots = self.lu_solve_test_helper(A_dims, b_dims, pivot, device, dtype)
                x_exp_list = []
                for i in range(b_dims[0]):
                    x_exp_list.append(torch.lu_solve(b[i], LU_data[i], LU_pivots[i]))
                x_exp = torch.stack(x_exp_list)  # Stacked output
                x_act = torch.lu_solve(b, LU_data, LU_pivots)  # Actual output
                self.assertEqual(x_exp, x_act)  # Equality check
                Ax = np.matmul(A.cpu(), x_act.cpu())
                self.assertEqual(b, Ax)

            for batchsize in [1, 3, 4]:
                lu_solve_batch_test_helper((batchsize, 5, 5), (batchsize, 5, 10), pivot)

        # Tests tensors with 0 elements
        b = torch.randn(3, 0, 3, dtype=dtype, device=device)
        A = torch.randn(3, 0, 0, dtype=dtype, device=device)
        LU_data, LU_pivots = torch.lu(A)
        self.assertEqual(torch.empty_like(b), b.lu_solve(LU_data, LU_pivots))

        sub_test(True)
        if self.device_type == 'cuda':
            sub_test(False)
//  [python]
def test_lu_solve_batched(self, device, dtype):
        def sub_test(pivot):
            def lu_solve_batch_test_helper(A_dims, b_dims, pivot):
                b, A, LU_data, LU_pivots = self.lu_solve_test_helper(A_dims, b_dims, pivot, device, dtype)
                x_exp_list = []
                for i in range(b_dims[0]):
                    x_exp_list.append(torch.linalg.solve(b[i], torch.linalg.lu_factor(A[i], pivot=pivot)))
                x_exp = torch.stack(x_exp_list)  # Stacked output
                x_act = torch.linalg.solve(b, torch.linalg.lu_factor(A, pivot=pivot))  # Actual output
                self.assertEqual(x_exp, x_act)  # Equality check
                Ax = np.matmul(A.cpu(), x_act.cpu())
                self.assertEqual(b, Ax)

            for batchsize in [1, 3, 4]:
                lu_solve_batch_test_helper((batchsize, 5, 5), (batchsize, 5, 10), pivot)

        # Tests tensors with 0 elements
        b = torch.randn(3, 0, 3, dtype=dtype, device=device)
        A = torch.randn(3, 0, 0, dtype=dtype, device=device)
        LU_data, LU_pivots = torch.linalg.lu_factor(A)
        self.assertEqual(torch.empty_like(b), torch.linalg.solve(b, LU_data, LU_pivots))

        sub_test(True)
        if self.device_type == 'cuda':
            sub_test(False)
//  [python]
def __init__(self, in_features):
        super().__init__()
        self.in_features = in_features

        weight = torch.randn(self.in_features, self.in_features)
        weight = torch.nn.init.orthogonal_(weight)
        d = torch.lu(weight)
        w_p, w_l, w_u = torch.lu_unpack(d[0], d[1])
        w_s = torch.diag(w_u)
        w_u = torch.triu(w_u, 1)
        u_mask = torch.triu(torch.ones_like(w_u), 1)
        l_mask = u_mask.T

        self.register_buffer("w_p", w_p)
        self.register_buffer("u_mask", u_mask)
        self.register_buffer("l_mask", l_mask)
        self.register_buffer("s_sign", torch.sign(w_s))
        self.register_buffer("l_eye", torch.eye(l_mask.shape[0]))
        self.w_l = torch.nn.Parameter(w_l)
        self.w_s = torch.nn.Parameter(logabs(w_s))
        self.w_u = torch.nn.Parameter(w_u)
//  [python]
def __init__(self, in_features):
        super().__init__()
        self.in_features = in_features

        weight = torch.randn(self.in_features, self.in_features)
        weight = torch.nn.init.orthogonal_(weight)
        d = torch.linalg.lu_factor(weight)
        w_p, w_l, w_u = torch.linalg.lu_unpack(*d)
        w_s = torch.diag(w_u)
        w_u = torch.triu(w_u, 1)
        u_mask = torch.triu(torch.ones_like(w_u), 1)
        l_mask = u_mask.T

        self.register_buffer("w_p", w_p)
        self.register_buffer("u_mask", u_mask)
        self.register_buffer("l_mask", l_mask)
        self.register_buffer("s_sign", torch.sign(w_s))
        self.register_buffer("l_eye", torch.eye(l_mask.shape[0]))
        self.w_l = torch.nn.Parameter(w_l)
        self.w_s = torch.nn.Parameter(logabs(w_s))
        self.w_u = torch.nn.Parameter(w_u)
//  [python]
def setup(self):
        self.giHig = self.lambd * self.dots
        diag = torch.diag(torch.full([self.m], self.m, device=self.dev, dtype=self.dtype))
        #self.giHig = torch.lu(self.giHig + diag, pivot=False)[0]
        self.giHig = torch.lu(self.giHig + diag)[0]
        self.giHig = torch.triu(self.giHig - diag)
        self.denom = self.m + torch.diagonal(self.giHig)
        tmp = -self.giHig.t().contiguous() / self.denom.reshape((1, -1))

        if USE_CUDA:
            diag = torch.diag(torch.full([self.m], self.lambd, device=self.dev, dtype=self.dtype))
            self.coef = hinv_cuda.hinv_setup(tmp, diag)
        else:
            for i in range(max(self.last, 1), self.m):
                self.coef[i, :i] = tmp[i, :i].matmul(self.coef[:i, :i])
//  [python]
def setup(self):
        self.giHig = self.lambd * self.dots
        diag = torch.diag(torch.full([self.m], self.m, device=self.dev, dtype=self.dtype))
        self.giHig = torch.linalg.lu_factor(self.giHig + diag)[0]
        self.giHig = torch.triu(self.giHig - diag)
        self.denom = self.m + torch.diagonal(self.giHig)
        tmp = -self.giHig.t().contiguous() / self.denom.reshape((1, -1))

        if USE_CUDA:
            diag = torch.diag(torch.full([self.m], self.lambd, device=self.dev, dtype=self.dtype))
            self.coef = hinv_cuda.hinv_setup(tmp, diag)
        else:
            for i in range(max(self.last, 1), self.m):
                self.coef[i, :i] = tmp[i, :i].matmul(self.coef[:i, :i])
//  [python]
def rank_revealing_LUP_GPU(A, threshold=10**-10):
    """
    It returns the rank of A using GPU acceleration. This is based on PyTorch's LUP implementation.
    
    Parameters
    __________

    A: (2D float64 Numpy Array)
    Input array for which rank needs to be computed

    threshold: (Float)
    See documentation of the function "sort_rows_fast"


    Returns:
    _________

    rank: (Integer)
    rank of A

    """

    try:
        import torch
    except ImportError as e:
        print(str(e))
        print("rank_revealing_LUP_GPU requires PyTorch to run.")

    A = A / np.abs(A).max()
    A_tensor = torch.from_numpy(A).cuda()    
    A_LU, pivots = torch.lu(A_tensor, get_infos=False)
    _, _, u = torch.lu_unpack(A_LU, pivots, unpack_pivots=False)
    u = u.cpu().numpy()
    ref = compute_ref_LUP_fast(u, threshold)
    rank = compute_rank_ref(ref, threshold)

    return rank
//  [python]
def rank_revealing_LUP_GPU(A, threshold=10**-10):
    """
    It returns the rank of A using GPU acceleration. This is based on PyTorch's LUP implementation.
    
    Parameters
    __________

    A: (2D float64 Numpy Array)
    Input array for which rank needs to be computed

    threshold: (Float)
    See documentation of the function "sort_rows_fast"


    Returns:
    _________

    rank: (Integer)
    rank of A

    """

    try:
        import torch
    except ImportError as e:
        print(str(e))
        print("rank_revealing_LUP_GPU requires PyTorch to run.")

    A = A / np.abs(A).max()
    A_tensor = torch.from_numpy(A).cuda()    
    LU, pivots = torch.linalg.lu_factor(A_tensor)
    _, _, u = torch.lu_unpack(LU, pivots, unpack_pivots=False)
    u = u.cpu().numpy()
    ref = compute_ref_LUP_fast(u, threshold)
    rank = compute_rank_ref(ref, threshold)

    return rank
//  [python]
def solve(
        self,
        damping: Optional[Union[float, torch.Tensor]] = None,
        ellipsoidal_damping: bool = True,
        damping_eps: float = 1e-8,
        **kwargs,
    ) -> torch.Tensor:
        if self._check_singular:
            AtA = self.linearization.AtA
            Atb = self.linearization.Atb
            with torch.no_grad():
                output = torch.zeros(AtA.shape[0], AtA.shape[1]).to(AtA.device)
                _, _, infos = torch.lu(AtA, get_infos=True)
                good_idx = infos.bool().logical_not()
                if not good_idx.all():
                    warnings.warn(
                        "Singular matrix found in batch, solution will be set "
                        "to all 0 for all singular matrices.",
                        RuntimeWarning,
                    )
            AtA = AtA[good_idx]
            Atb = Atb[good_idx]
            solution = self._apply_damping_and_solve(
                Atb,
                AtA,
                damping=damping,
                ellipsoidal_damping=ellipsoidal_damping,
                damping_eps=damping_eps,
            )
            output[good_idx] = solution
            return output
        else:
            return self._apply_damping_and_solve(
                self.linearization.Atb,
                self.linearization.AtA,
                damping=damping,
                ellipsoidal_damping=ellipsoidal_damping,
                damping_eps=damping_eps,
            )
//  [python]
def solve(
        self,
        damping: Optional[Union[float, torch.Tensor]] = None,
        ellipsoidal_damping: bool = True,
        damping_eps: float = 1e-8,
        **kwargs,
    ) -> torch.Tensor:
        if self._check_singular:
            AtA = self.linearization.AtA
            Atb = self.linearization.Atb
            with torch.no_grad():
                output = torch.zeros(AtA.shape[0], AtA.shape[1]).to(AtA.device)
                LU, pivots = torch.linalg.lu_factor(AtA)
                infos = torch.linalg.lu_factor_ex(AtA).info
                good_idx = infos.bool().logical_not()
                if not good_idx.all():
                    warnings.warn(
                        "Singular matrix found in batch, solution will be set "
                        "to all 0 for all singular matrices.",
                        RuntimeWarning,
                    )
            AtA = AtA[good_idx]
            Atb = Atb[good_idx]
            solution = self._apply_damping_and_solve(
                Atb,
                AtA,
                damping=damping,
                ellipsoidal_damping=ellipsoidal_damping,
                damping_eps=damping_eps,
            )
            output[good_idx] = solution
            return output
        else:
            return self._apply_damping_and_solve(
                self.linearization.Atb,
                self.linearization.AtA,
                damping=damping,
                ellipsoidal_damping=ellipsoidal_damping,
                damping_eps=damping_eps,
            )
//  [python]
def lu_solve_test_helper(self, A_dims, b_dims, pivot, device, dtype):
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = random_fullrank_matrix_distinct_singular_value(*A_dims, dtype=dtype, device=device)
        LU_data, LU_pivots, info = torch.lu(A, get_infos=True, pivot=pivot)
        self.assertEqual(info, torch.zeros_like(info))
        return b, A, LU_data, LU_pivots
//  [python]
def lu_solve_test_helper(self, A_dims, b_dims, pivot, device, dtype):
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = random_fullrank_matrix_distinct_singular_value(*A_dims, dtype=dtype, device=device)
        LU_data, LU_pivots = torch.linalg.lu_factor(A, pivot=pivot)
        self.assertEqual(LU_pivots, torch.zeros_like(LU_pivots))
        return b, A, LU_data, LU_pivots
//  [python]
def forward(ctx, A, b):
        A_lu,pivots,check=torch.lu(A,get_infos=True)
        if check.abs().item()>1.e-5: #non-zero, fail:
            print('A matrix is singular1')
            A_lu,pivots,check=torch.lu(A+0.5*torch.eye(A.shape[0],dtype=A.dtype,device=A.device),get_infos=True)
            if check.abs().item()>1.e-5:
                print('A matrix is singular2')
                x=torch.zeros((b.shape[0], 1),device=b.device)
                A_lu.zero_()
                ctx.save_for_backward(A_lu,pivots,x)
                return x
        x=torch.lu_solve(b.view(-1,1),A_lu,pivots)
        # if (torch.isnan(x)).any().item() or (torch.isinf(x)).any().item():
        #     print('solution is nan')
        #     x=torch.zeros((b.shape[0], 1),device=b.device)
        #     A_lu.zero_()
        #     ctx.save_for_backward(A_lu,pivots,x)
        #     return x
        ctx.save_for_backward(A_lu,pivots,x)
        return x
//  [python]
def forward(ctx, A, b):
        A_lu,pivots=torch.linalg.lu_factor(A)
        check=torch.linalg.det(A)
        if check.abs().item()>1.e-5: #non-zero, fail:
            print('A matrix is singular1')
            A_lu,pivots=torch.linalg.lu_factor(A+0.5*torch.eye(A.shape[0],dtype=A.dtype,device=A.device))
            check=torch.linalg.det(A+0.5*torch.eye(A.shape[0],dtype=A.dtype,device=A.device))
            if check.abs().item()>1.e-5:
                print('A matrix is singular2')
                x=torch.zeros((b.shape[0], 1),device=b.device)
                A_lu.zero_()
                ctx.save_for_backward(A_lu,pivots,x)
                return x
        x=torch.linalg.solve(A_lu,b.view(-1,1),pivots)
        ctx.save_for_backward(A_lu,pivots,x)
        return x
//  [python]
def forward(ctx, A, b, Graph_extend_batch):
        x=[]
        A_lu=[]
        for bid in range(Graph_extend_batch[-1].item()+1):
            mask=Graph_extend_batch==bid
            Mask=mask[:,None]&mask[None,:]
            Ab=A[Mask].reshape(-1,mask.sum().item())           
            bb=b[mask]
            Ab_lu,Ab_pivots,check=torch.lu(Ab,get_infos=True)
            if check.abs().item()>1.e-5:
                Ab_lu,Ab_pivots,check=torch.lu(Ab+0.5*torch.eye(Ab.shape[0],dtype=Ab.dtype,device=Ab.device),get_infos=True)
                if check.abs().item()>1.e-5:
                    xb=torch.zeros_like(bb).view(-1,1)
                    Ab_lu.zero_()
                    x.append(xb)
                    A_lu.extend([Ab_lu,Ab_pivots])
                    continue            
            xb=torch.lu_solve(bb.view(-1,1),Ab_lu,Ab_pivots)            
            x.append(xb)
            A_lu.extend([Ab_lu,Ab_pivots])        
        ctx.save_for_backward(*A_lu,*x,Graph_extend_batch)
        x=torch.cat(x,dim=0)
        return x
//  [python]
def forward(ctx, A, b, Graph_extend_batch):
        x=[]
        A_lu=[]
        for bid in range(Graph_extend_batch[-1].item()+1):
            mask=Graph_extend_batch==bid
            Mask=mask[:,None]&mask[None,:]
            Ab=A[Mask].reshape(-1,mask.sum().item())           
            bb=b[mask]
            LU,pivots=torch.linalg.lu_factor(Ab)
            check=torch.linalg.det(LU)
            if check.abs().item()>1.e-5:
                LU,pivots=torch.linalg.lu_factor(Ab+0.5*torch.eye(Ab.shape[0],dtype=Ab.dtype,device=Ab.device))
                check=torch.linalg.det(LU)
                if check.abs().item()>1.e-5:
                    xb=torch.zeros_like(bb).view(-1,1)
                    LU.zero_()
                    x.append(xb)
                    A_lu.extend([LU,pivots])
                    continue            
            xb=torch.linalg.solve(LU,bb.view(-1,1),pivots)            
            x.append(xb)
            A_lu.extend([LU,pivots])        
        ctx.save_for_backward(*A_lu,*x,Graph_extend_batch)
        x=torch.cat(x,dim=0)
        return x
//  [python]
def batch_det(A):
    """
    (No use) Batch compute determinant of square matrix A of shape (*, N, N)

    We can use torch.linalg.det() directly
    Return:
    Tensor of shape (*)

    第一种初等变换—bai—交换两行或du列—zhi—要偶数次不改变dao行列的值
    第二种初等zhuan变换——shu某行（列）乘以非0实数——这个可以乘以系数，但总的乘积必须为1方可不改变行列式值
    第三种初等变换——某行（列）乘以实数加到另一行（列）上——此条对行列式值无影响
    """
    # 对矩阵进行LU 分解，显然，分解后的矩阵对角线上元素的乘积即为原始矩阵行列式的值
    # 由于在进行 LU 分解时可能会进行行交换的情况，而每次行交换都会带来行列式的符号变化，所以我们要记录行交换次数的奇偶性，
    # 奇数则符号改变，偶数则符号不变，请密切注意下面程序中出现的用来标记变换次数奇偶性的变量 parity 
    # LU: the packed LU factorization matrix, (*, N, N)
    # pivots: (*, N)
    LU, pivots = torch.lu(A)
    # torch.einsum('...ii->...i', LU): the diagnoal vector
    # det_LU: the product of all diagnal values
    det_LU = torch.einsum('...ii->...i', LU).prod(-1)
    pivots -= 1
    d = pivots.shape[-1]
    perm = pivots - torch.arange(d, dtype=pivots.dtype, device=pivots.device).expand(pivots.shape)
    det_P = (-1) ** ((perm != 0).sum(-1))
    det = det_LU * det_P.type(det_LU.dtype)

    return det
//  [python]
def batch_det(A):
    """
    (No use) Batch compute determinant of square matrix A of shape (*, N, N)

    We can use torch.linalg.det() directly
    Return:
    Tensor of shape (*)

    第一种初等变换—bai—交换两行或du列—zhi—要偶数次不改变dao行列的值
    第二种初等zhuan变换——shu某行（列）乘以非0实数——这个可以乘以系数，但总的乘积必须为1方可不改变行列式值
    第三种初等变换——某行（列）乘以实数加到另一行（列）上——此条对行列式值无影响
    """
    # 对矩阵进行LU 分解，显然，分解后的矩阵对角线上元素的乘积即为原始矩阵行列式的值
    # 由于在进行 LU 分解时可能会进行行交换的情况，而每次行交换都会带来行列式的符号变化，所以我们要记录行交换次数的奇偶性，
    # 奇数则符号改变，偶数则符号不变，请密切注意下面程序中出现的用来标记变换次数奇偶性的变量 parity 
    # LU: the packed LU factorization matrix, (*, N, N)
    # pivots: (*, N)
    LU, pivots = torch.linalg.lu_factor(A)
    # torch.einsum('...ii->...i', LU): the diagnoal vector
    # det_LU: the product of all diagnal values
    det_LU = torch.einsum('...ii->...i', LU).prod(-1)
    pivots -= 1
    d = pivots.shape[-1]
    perm = pivots - torch.arange(d, dtype=pivots.dtype, device=pivots.device).expand(pivots.shape)
    det_P = (-1) ** ((perm != 0).sum(-1))
    det = det_LU * det_P.type(det_LU.dtype)

    return det
//  [python]
def batch_det(A):
    """
    Batch determinant of square matrix A of shape (*, N, N)
    Return:
    Tensor of shape (*)
    """
    LU, pivots = torch.lu(A)
    det_LU = torch.einsum('...ii->...i', LU).prod(-1)
    pivots -= 1
    d = pivots.shape[-1]
    perm = pivots - torch.arange(d, dtype=pivots.dtype, device=pivots.device).expand(pivots.shape)
    det_P = (-1) ** ((perm != 0).sum(-1))
    det = det_LU * det_P.type(det_LU.dtype)

    return det
//  [python]
def batch_det(A):
    """
    Batch determinant of square matrix A of shape (*, N, N)
    Return:
    Tensor of shape (*)
    """
    LU, pivots = torch.linalg.lu_factor(A)
    det_LU = torch.einsum('...ii->...i', LU).prod(-1)
    pivots -= 1
    d = pivots.shape[-1]
    perm = pivots - torch.arange(d, dtype=pivots.dtype, device=pivots.device).expand(pivots.shape)
    det_P = (-1) ** ((perm != 0).sum(-1))
    det = det_LU * det_P.type(det_LU.dtype)

    return det
//  [python]
def __init__(self, likelihood:MultitaskMultivariateNormal, train_targets, USE_LU=True):
        self.USE_LU = USE_LU # if set to False, use Cholesky factorization
        err = (train_targets - likelihood.mean).view(-1,1).contiguous()
        if self.USE_LU:
            self.LU = torch.lu(likelihood.covariance_matrix)
            self.alpha = torch.lu_solve(err, *self.LU )
        else:  
            self.L = torch.cholesky(likelihood.covariance_matrix, upper=False)
            self.alpha = torch.triangular_solve( torch.triangular_solve(err, self.L, upper=False).solution, self.L.T, upper=True).solution
//  [python]
def __init__(self, likelihood:MultitaskMultivariateNormal, train_targets, USE_LU=True):
        self.USE_LU = USE_LU # if set to False, use Cholesky factorization
        err = (train_targets - likelihood.mean).view(-1,1).contiguous()
        if self.USE_LU:
            self.LU = torch.linalg.lu_factor(likelihood.covariance_matrix)
            self.alpha = torch.linalg.lu_solve(err, *self.LU )
        else:  
            self.L = torch.cholesky(likelihood.covariance_matrix, upper=False)
            self.alpha = torch.triangular_solve( torch.triangular_solve(err, self.L, upper=False).solution, self.L.T, upper=True).solution
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^3 + D^2N)
            logabsdet = O(D^3)
        where:
            D = num of features
            N = num of inputs
        """
        batch_size = inputs.shape[0]
        outputs = inputs - self.bias
        lu, lu_pivots = torch.lu(self._weight)
        outputs = torch.lu_solve(outputs.t(), lu, lu_pivots).t()
        logabsdet = -torch.sum(torch.log(torch.abs(torch.diag(lu))))
        logabsdet = logabsdet * outputs.new_ones(batch_size)
        return outputs, logabsdet