jiahuimbzuai/CoQuIR huggingface dataset:
//  [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.linalg.lu_factor(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
//  [python]
def fit(self):
        """
        Fit polynomial to data.
        """
        import networkx as nx
        import torch
        # Step 1. Calculate the Laplacian matrix
        L = nx.laplacian_matrix(self.Graph)
        nodelist = self.Graph.nodes()
        K = L.shape[0]

        # Step 2. Get the data in the right format 
        cache = self.loss_function(self.data_train)
        
        # Step 3. Compute the proximal loss
        def proximal_loss(t, nu, warm_start, pool, cache=cache):
            XtX = cache['XtX']
            XtY = cache['XtY']
            n = cache['n']
            # LU = X'X + 0.5 * t * I
            Alu = torch.lu(XtX + 1./(2 * t) * torch.eye(n).unsqueeze(0).double())
            b = XtY + 1./(2 * t) * torch.from_numpy(nu)
            x = torch.lu_solve(b, *Alu).numpy()
            return x

        def proximal_residual(t, nu, warm_start, pool, lambda_val=1e-4):
            return nu / (1. + t * lambda_val)

        G_to_data = self._graph_to_data(cache['alpha_shape'])
        result, info = self._stratified_model_admm(shape=cache['shape'], \
                                            Lap=L, \
                                            loss_proximal_func=proximal_loss, \
                                            regulariser_proximal_func=proximal_residual, \
                                            graph_data=G_to_data)
        print(info)
        return self._output_to_graph(result)
//  [python]
def fit(self):
        """
        Fit polynomial to data.
        """
        import networkx as nx
        import torch
        # Step 1. Calculate the Laplacian matrix
        L = nx.laplacian_matrix(self.Graph)
        nodelist = self.Graph.nodes()
        K = L.shape[0]

        # Step 2. Get the data in the right format 
        cache = self.loss_function(self.data_train)
        
        # Step 3. Compute the proximal loss
        def proximal_loss(t, nu, warm_start, pool, cache=cache):
            XtX = cache['XtX']
            XtY = cache['XtY']
            n = cache['n']
            # LU = X'X + 0.5 * t * I
            Alu = torch.linalg.lu_factor(XtX + 1./(2 * t) * torch.eye(n).unsqueeze(0).double())
            b = XtY + 1./(2 * t) * torch.from_numpy(nu)
            x = torch.linalg.lu_solve(b, *Alu).numpy()
            return x

        def proximal_residual(t, nu, warm_start, pool, lambda_val=1e-4):
            return nu / (1. + t * lambda_val)

        G_to_data = self._graph_to_data(cache['alpha_shape'])
        result, info = self._stratified_model_admm(shape=cache['shape'], \
                                            Lap=L, \
                                            loss_proximal_func=proximal_loss, \
                                            regulariser_proximal_func=proximal_residual, \
                                            graph_data=G_to_data)
        print(info)
        return self._output_to_graph(result)
//  [python]
def weight_inverse_and_logabsdet(self):
        """
        Cost:
            inverse = O(D^3)
            logabsdet = O(D)
        where:
            D = num of features
        """
        identity = torch.eye(self.features, self.features)
        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
        """
        identity = torch.eye(self.features, self.features)
        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_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
                self.assertLessEqual(b.dist(torch.matmul(A, x_act)), 1e-12)  # Correctness check

            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_triangular(b[i], LU_data[i], upper=False, unitriangular=False, left=True))
                x_exp = torch.stack(x_exp_list)  # Stacked output
                x_act = torch.linalg.solve_triangular(b, LU_data, upper=False, unitriangular=False, left=True)  # Actual output
                self.assertEqual(x_exp, x_act)  # Equality check
                self.assertLessEqual(b.dist(torch.matmul(A, x_act)), 1e-12)  # Correctness check

            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_triangular(b, LU_data, upper=False, unitriangular=False, left=True))

        sub_test(True)
        if self.device_type == 'cuda':
            sub_test(False)
//  [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-decompose the weights and solve for the outputs.
        lu, lu_pivots = torch.lu(self._weight)
        outputs = torch.lu_solve(outputs.t(), lu, lu_pivots).t()
        # The linear-system solver returns the LU decomposition of the weights, which we
        # can use to obtain the log absolute determinant directly.
        logabsdet = -torch.sum(torch.log(torch.abs(torch.diag(lu))))
        logabsdet = logabsdet * outputs.new_ones(batch_size)
        return outputs, logabsdet
//  [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-decompose the weights and solve for the outputs.
        lu, pivots = torch.linalg.lu_factor(self._weight)
        outputs = torch.linalg.lu_solve(outputs.t(), lu, pivots).t()
        # The linear-system solver returns the LU decomposition of the weights, which we
        # can use to obtain the log absolute determinant directly.
        logabsdet = -torch.sum(torch.log(torch.abs(torch.diag(lu))))
        logabsdet = logabsdet * outputs.new_ones(batch_size)
        return outputs, logabsdet
//  [python]
def Solver(A, b, method="cholesky"):
        # methods: 1) 'lu': LU linear solver.
        # 2) "cholesky": Cholesky Solver.

        # Code ref: https://github.com/DeformableFriends/NeuralTracking/blob/main/model/model.py
        if method == "lu":
            A_LU, pivots = torch.lu(A)
            x = torch.lu_solve(b, A_LU, pivots)

        elif method == "cholesky":
            U = torch.linalg.cholesky(A)
            x = torch.cholesky_solve(b, U)
        
        return x
//  [python]
def Solver(A, b, method="cholesky"):
        # methods: 1) 'lu': LU linear solver.
        # 2) "cholesky": Cholesky Solver.

        # Code ref: https://github.com/DeformableFriends/NeuralTracking/blob/main/model/model.py
        if method == "lu":
            LU, pivots = torch.linalg.lu_factor(A)
            x = torch.linalg.lu_solve(b, LU, pivots)

        elif method == "cholesky":
            U = torch.linalg.cholesky(A)
            x = torch.cholesky_solve(b, U)
        
        return x
//  [python]
def forward(ctx, A: torch.Tensor, b: torch.Tensor):
        l, u = torch.lu(A)
        d = torch.lu_solve(b, l, u)  # D = Q-1 @ P
        ctx.save_for_backward(l, u, d)
        return d
//  [python]
def forward(ctx, A: torch.Tensor, b: torch.Tensor):
        LU, pivots = torch.linalg.lu_factor(A)
        d = torch.linalg.lu_solve(b, LU, pivots)  # D = Q-1 @ P
        ctx.save_for_backward(LU, pivots, d)
        return d
//  [python]
def prox(self, t, nu, warm_start, pool, cache):
		# raise NotImplementedError("This method is not yet done!!!")

		XtX = cache['XtX']
		XtY = cache['XtY']
		n = cache['n']

		A_LU = torch.lu(
			XtX + 1. / (2 * t) * torch.eye(n).unsqueeze(0).double())
		b = XtY + 1. / (2 * t) * torch.from_numpy(nu)
		x = torch.lu_solve(b, *A_LU)

		return x.numpy()
//  [python]
def prox(self, t, nu, warm_start, pool, cache):
		XtX = cache['XtX']
		XtY = cache['XtY']
		n = cache['n']

		A_LU = torch.linalg.lu_factor(
			XtX + 1. / (2 * t) * torch.eye(n).unsqueeze(0).double())
		b = XtY + 1. / (2 * t) * torch.from_numpy(nu)
		x = torch.linalg.lu_solve(b, *A_LU)

		return x.numpy()
//  [python]
def __init__(self, dim):
        super(Invertible1x1Conv, self).__init__()
        self.dim = dim

        # Grab the weight and bias from a randomly initialized Conv2d.
        m = nn.Conv2d(dim, dim, kernel_size=1)
        W = m.weight.clone().detach().reshape(dim, dim)
        LU, pivots = torch.lu(W)
        P, _, _ = torch.lu_unpack(LU, pivots)

        s = torch.diag(LU)
        # noinspection PyTypeChecker
        LU = torch.where(torch.eye(dim) == 0, LU, torch.zeros_like(LU))

        self.register_buffer("P", P)
        self.register_buffer("s_sign", torch.sign(s))
        self.register_parameter("s_log", nn.Parameter(torch.log(torch.abs(s) + 1e-3)))
        self.register_parameter("LU", nn.Parameter(LU))
//  [python]
def __init__(self, dim):
        super(Invertible1x1Conv, self).__init__()
        self.dim = dim

        # Grab the weight and bias from a randomly initialized Conv2d.
        m = nn.Conv2d(dim, dim, kernel_size=1)
        W = m.weight.clone().detach().reshape(dim, dim)
        LU, pivots = torch.linalg.lu_factor(W)
        P, _, _ = torch.lu_unpack(LU, pivots)

        s = torch.diag(LU)
        # noinspection PyTypeChecker
        LU = torch.where(torch.eye(dim) == 0, LU, torch.zeros_like(LU))

        self.register_buffer("P", P)
        self.register_buffer("s_sign", torch.sign(s))
        self.register_parameter("s_log", nn.Parameter(torch.log(torch.abs(s) + 1e-3)))
        self.register_parameter("LU", nn.Parameter(LU))
//  [python]
def lu_solve_test_helper(self, A_dims, b_dims, pivot, device, dtype):
        from 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 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 test_lu_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from 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 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 sampling_W(self, dim):
        # sample a rotation matrix
        W = torch.empty(dim, dim)
        torch.nn.init.orthogonal_(W)
        # compute LU
        LU, pivot = torch.lu(W)
        P, L, U = torch.lu_unpack(LU, pivot)
        return W, P, L, U
//  [python]
def sampling_W(self, dim):
        # sample a rotation matrix
        W = torch.empty(dim, dim)
        torch.nn.init.orthogonal_(W)
        # compute LU
        LU, pivots = torch.linalg.lu_factor(W)
        P, L, U = torch.lu_unpack(LU, pivots)
        return W, P, L, U
//  [python]
def sample_inputs_lu_unpack(op_info, device, dtype, requires_grad=False, **kwargs):
    for lu_sample in sample_inputs_lu(op_info, device, dtype, requires_grad, **kwargs):
        lu_data, pivots = torch.lu(lu_sample.input)
        lu_data.requires_grad_(requires_grad)
        yield SampleInput(lu_data, args=(pivots,))
//  [python]
def sample_inputs_lu_unpack(op_info, device, dtype, requires_grad=False, **kwargs):
    for lu_sample in sample_inputs_lu(op_info, device, dtype, requires_grad, **kwargs):
        lu_data, pivots = torch.linalg.lu_factor(lu_sample.input)
        lu_data.requires_grad_(requires_grad)
        yield SampleInput(lu_data, args=(pivots,))
//  [python]
def test_lu_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve
        from 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 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 sample_inputs_lu_unpack(op_info, device, dtype, requires_grad=False, **kwargs):
    def out_fn(output):
        return output[1], output[2]

    for lu_sample in sample_inputs_linalg_lu(op_info, device, dtype, requires_grad, **kwargs):
        lu_data, pivots = torch.lu_unpack(*torch.lu(lu_sample.input))
        lu_data.requires_grad_(requires_grad)
        yield SampleInput(lu_data, args=(pivots,), output_process_fn_grad=out_fn)
//  [python]
def sample_inputs_lu_unpack(op_info, device, dtype, requires_grad=False, **kwargs):
    def out_fn(output):
        return output[1], output[2]

    for lu_sample in sample_inputs_linalg_lu(op_info, device, dtype, requires_grad, **kwargs):
        lu_data, pivots = torch.linalg.lu_factor(lu_sample.input)
        lu_data.requires_grad_(requires_grad)
        yield SampleInput(lu_data, args=(pivots,), output_process_fn_grad=out_fn)
//  [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.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, pivot=False)[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 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_ge(1, 10):
        sol = _torch_solve_cast(A, B)
        warnings.warn("PyTorch version < 1.10, solve validness mask maybe not correct", RuntimeWarning)
        return sol, sol, 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: torch.dtype = B.dtype
    if dtype not in (torch.float32, torch.float64):
        dtype = torch.float32

    if TYPE_CHECKING:
        # TODO: remove this branch when kornia relies on torch >= 1.13
        A_LU: Tensor
        pivots: Tensor
        info: Tensor
    elif torch_version_ge(1, 13):
        A_LU, pivots, info = torch.linalg.lu_factor_ex(A.to(dtype))
    else:
        # TODO: remove this branch when kornia relies on torch >= 1.13
        A_LU, pivots, info = torch.lu(A.to(dtype), True, get_infos=True)

    valid_mask: Tensor = info == 0
    n_dim_B = len(B.shape)
    n_dim_A = len(A.shape)
    if n_dim_A - n_dim_B == 1:
        B = B.unsqueeze(-1)

    if TYPE_CHECKING:
        # TODO: remove this branch when kornia relies on torch >= 1.13
        X: Tensor
    elif torch_version_ge(1, 13):
        X = torch.linalg.lu_solve(A_LU, pivots, B.to(dtype))
    else:
        # TODO: remove this branch when kornia relies on torch >= 1.13
        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_ge(1, 10):
        sol = _torch_solve_cast(A, B)
        warnings.warn("PyTorch version < 1.10, solve validness mask maybe not correct", RuntimeWarning)
        return sol, sol, 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: torch.dtype = B.dtype
    if dtype not in (torch.float32, torch.float64):
        dtype = torch.float32

    if TYPE_CHECKING:
        # TODO: remove this branch when kornia relies on torch >= 1.13
        A_LU: Tensor
        pivots: Tensor
        info: Tensor
    else:
        A_LU, pivots, info = torch.linalg.lu_factor_ex(A.to(dtype))

    valid_mask: Tensor = info == 0
    n_dim_B = len(B.shape)
    n_dim_A = len(A.shape)
    if n_dim_A - n_dim_B == 1:
        B = B.unsqueeze(-1)

    if TYPE_CHECKING:
        # TODO: remove this branch when kornia relies on torch >= 1.13
        X: Tensor
    else:
        X = torch.linalg.lu_solve(A_LU, pivots, B.to(dtype))

    return X.to(B.dtype), A_LU.to(A.dtype), valid_mask
//  [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-decompose the weights and solve for the outputs.
        lu, pivots = torch.lu_unpack(*torch.lu(self._weight))
        outputs = torch.lu_solve(lu, outputs.t(), pivots).t()
        # The linear-system solver returns the LU decomposition of the weights, which we
        # can use to obtain the log absolute determinant directly.
        logabsdet = -torch.sum(torch.log(torch.abs(torch.diag(lu))))
        logabsdet = logabsdet * outputs.new_ones(batch_size)
        return outputs, logabsdet
//  [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-decompose the weights and solve for the outputs.
        lu, lu_pivots = torch.linalg.lu_factor(self._weight)
        outputs = torch.linalg.lu_solve(lu, lu_pivots, outputs.t()).t()
        # The linear-system solver returns the LU decomposition of the weights, which we
        # can use to obtain the log absolute determinant directly.
        logabsdet = -torch.sum(torch.log(torch.abs(torch.diag(lu))))
        logabsdet = logabsdet * outputs.new_ones(batch_size)
        return outputs, logabsdet
//  [python]
def __init__(self, A):
        self.LU, self.pivots = torch.lu_unpack(*torch.lu(A))
//  [python]
def __init__(self, A):
        self.LU, self.pivots = torch.linalg.lu_factor(A)
//  [python]
def lu_factor(x):
    LU, pivots = torch.lu_unpack(*torch.lu(x))
    # torch retuns pivots with 1-based indexing
    return LU, pivots - 1
//  [python]
def lu_factor(x):
    LU, pivots = torch.linalg.lu_factor(x)
    # torch retuns pivots with 1-based indexing
    return LU, pivots - 1
//  [python]
def lu_factor(
    x: torch.Tensor,
    /,
    *,
    pivot: Optional[bool] = True,
    out: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    ret = torch.lu(x, pivot=pivot, out=out)
    ret_tuple = namedtuple("lu_factor", ["LU", "p"])
    return ret_tuple(ret[0], ret[1])
//  [python]
def lu_factor(
    x: torch.Tensor,
    /,
    *,
    pivot: Optional[bool] = True,
    out: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    ret = torch.linalg.lu_factor(x, pivot=pivot, out=out)
    ret_tuple = namedtuple("lu_factor", ["LU", "p"])
    return ret_tuple(ret.LU, ret.pivots)
//  [python]
def get_factor(self):
        if self._LU is None:
            A_tensor = dlpack.from_dlpack(self._A)
            self._LU, self._pivots = torch.lu_unpack(*torch.lu(A_tensor))
        return self._LU, self._pivots
//  [python]
def get_factor(self):
        if self._LU is None:
            A_tensor = dlpack.from_dlpack(self._A)
            self._LU, self._pivots = torch.linalg.lu_factor(A_tensor)
        return self._LU, self._pivots
//  [python]
def _kxyx_matmul(self, X, ny_pts_x, Y, ny_pts_y, rhs, **saved_data) -> Tuple[torch.Tensor, dict]:
        ker = self.kernel
        rhs = ker.mmv(X, ny_pts_x, rhs)
        rhs = ker.mmv(ny_pts_y, Y, rhs)
        if 'kymm_ludec' not in saved_data:
            ky_mm = id_regularize(
                ker(ny_pts_y, ny_pts_y),
                ker.params.cg_epsilon(Y.dtype) * ny_pts_y.shape[0]
            )
            saved_data['kymm_ludec'] = torch.lu_factor(ky_mm)
        LU, pivots = saved_data['kymm_ludec']
        rhs = torch.lu_solve(LU, pivots, rhs).contiguous()
        rhs = ker.mmv(Y, ny_pts_y, rhs)
        rhs = ker.mmv(ny_pts_x, X, rhs)
        return rhs, saved_data
//  [python]
def _kxyx_matmul(self, X, ny_pts_x, Y, ny_pts_y, rhs, **saved_data) -> Tuple[torch.Tensor, dict]:
        ker = self.kernel
        rhs = ker.mmv(X, ny_pts_x, rhs)
        rhs = ker.mmv(ny_pts_y, Y, rhs)
        if 'kymm_ludec' not in saved_data:
            ky_mm = id_regularize(
                ker(ny_pts_y, ny_pts_y),
                ker.params.cg_epsilon(Y.dtype) * ny_pts_y.shape[0]
            )
            saved_data['kymm_ludec'] = torch.linalg.lu_factor(ky_mm)
        LU, pivots = saved_data['kymm_ludec']
        rhs = torch.linalg.lu_solve(LU, pivots, rhs).contiguous()
        rhs = ker.mmv(Y, ny_pts_y, rhs)
        rhs = ker.mmv(ny_pts_x, X, rhs)
        return rhs, saved_data
//  [python]
def _solve(self, X, ny_pts_x, rhs, penalty, **saved_data) -> Tuple[torch.Tensor, dict]:
        if self.linsys_solver == 'lu':
            if 'gl_ludec' not in saved_data:
                gkx = self._explicit_glambda(X, ny_pts_x, penalty, 1e-3)
                LU, pivots = torch.lu_unpack(*torch.lu(gkx))
                saved_data['gl_ludec'] = (LU, pivots)
            LU, pivots = saved_data['gl_ludec']
            return torch.solve(rhs, LU, pivots).contiguous(), saved_data
        elif self.linsys_solver == 'cholesky':
            if 'gl_cho' not in saved_data:
                gkx = self._explicit_glambda(X, ny_pts_x, penalty, 1e-3)
                L, info = torch.linalg.cholesky_ex(gkx)
                if info > 0:
                    raise RuntimeError(f"Failed cholesky decomposition at leading minor of order {info.item()}")
                saved_data['gl_cho'] = L
            L = saved_data['gl_cho']
            return torch.cholesky_solve(rhs, L).contiguous(), saved_data
        elif self.linsys_solver == 'cg':
            cg = ConjugateGradient()
            out = cg.solve(None, rhs, self._implicit_glambda(X, ny_pts_x, penalty, 1e-5), 100)
            return out, saved_data
        else:
            raise ValueError(self.linsys_solver)
//  [python]
def _solve(self, X, ny_pts_x, rhs, penalty, **saved_data) -> Tuple[torch.Tensor, dict]:
        if self.linsys_solver == 'lu':
            if 'gl_ludec' not in saved_data:
                gkx = self._explicit_glambda(X, ny_pts_x, penalty, 1e-3)
                LU, pivots = torch.linalg.lu_factor(gkx)
                saved_data['gl_ludec'] = (LU, pivots)
            LU, pivots = saved_data['gl_ludec']
            return torch.linalg.lu_solve(LU, pivots, rhs).contiguous(), saved_data
        elif self.linsys_solver == 'cholesky':
            if 'gl_cho' not in saved_data:
                gkx = self._explicit_glambda(X, ny_pts_x, penalty, 1e-3)
                L, info = torch.linalg.cholesky_ex(gkx)
                if info > 0:
                    raise RuntimeError(f"Failed cholesky decomposition at leading minor of order {info.item()}")
                saved_data['gl_cho'] = L
            L = saved_data['gl_cho']
            return torch.cholesky_solve(rhs, L).contiguous(), saved_data
        elif self.linsys_solver == 'cg':
            cg = ConjugateGradient()
            out = cg.solve(None, rhs, self._implicit_glambda(X, ny_pts_x, penalty, 1e-5), 100)
            return out, saved_data
        else:
            raise ValueError(self.linsys_solver)
//  [python]
def fit(self, X, Y):
        """
        Args:
            X: (..., N, N)
            Y: (..., N, K)
        """
        LU, piv = torch.lu_unpack(*torch.lu(X.transpose(-1, -2) + torch.eye(X.shape[-2], device=X.device) * self.l))
        self.M = nn.Parameter(torch.lu_solve(Y, LU, piv).transpose(-1, -2))
//  [python]
def fit(self, X, Y):
        """
        Args:
            X: (..., N, N)
            Y: (..., N, K)
        """
        LU, piv = torch.linalg.lu_factor(X.transpose(-1, -2) + torch.eye(X.shape[-2], device=X.device) * self.l)
        self.M = nn.Parameter(torch.lu_solve(Y, LU, piv).transpose(-1, -2))
//  [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)

        with self.assertRaisesRegex(RuntimeError, "torch.int32 dtype"):
            torch.lu_unpack(lu_data, lu_pivots.long())

        # 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.numel() == 0 and u.numel() == 0)
        p, l, u = torch.lu_unpack(lu_data, lu_pivots, unpack_pivots=False)
        self.assertTrue(p.numel() == 0)
        p, l, u = torch.lu_unpack(lu_data, lu_pivots, unpack_data=False, unpack_pivots=False)
        self.assertTrue(p.numel() == 0 and l.numel() == 0 and u.numel() == 0)
//  [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)

        with self.assertRaisesRegex(RuntimeError, "torch.int32 dtype"):
            torch.lu_unpack(lu_data, lu_pivots.long())

        # 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.numel() == 0 and u.numel() == 0)
        p, l, u = torch.lu_unpack(lu_data, lu_pivots, unpack_pivots=False)
        self.assertTrue(p.numel() == 0)
        p, l, u = torch.lu_unpack(lu_data, lu_pivots, unpack_data=False, unpack_pivots=False)
        self.assertTrue(p.numel() == 0 and l.numel() == 0 and u.numel() == 0)
//  [python]
def test_linalg_lu_solve(self, device, dtype):
        make_arg = partial(make_tensor, dtype=dtype, device=device)

        backends = ["default"]

        if torch.device(device).type == 'cuda':
            if torch.cuda.has_magma:
                backends.append("magma")
            if has_cusolver():
                backends.append("cusolver")

        def gen_matrices():
            rhs = 3
            ns = (5, 2, 0)
            batches = ((), (0,), (1,), (2,), (2, 1), (0, 2))
            for batch, n in product(batches, ns):
                yield make_arg(batch + (n, n)), make_arg(batch + (n, rhs))
            # Shapes to exercise all the paths
            shapes = ((1, 64), (2, 128), (1025, 2))
            for b, n in shapes:
                yield make_arg((b, n, n)), make_arg((b, n, rhs))


        for A, B in gen_matrices():
            LU, pivots = torch.lu(A)
            for backend in backends:
                torch.backends.cuda.preferred_linalg_library(backend)

                for left, adjoint in product((True, False), repeat=2):
                    B_left = B if left else B.mT
                    X = torch.solve(B_left, LU, pivots, left=left, adjoint=adjoint)
                    A_adj = A.mH if adjoint else A
                    if left:
                        self.assertEqual(B_left, A_adj @ X)
                    else:
                        self.assertEqual(B_left, X @ A_adj)
//  [python]
def test_linalg_lu_solve(self, device, dtype):
        make_arg = partial(make_tensor, dtype=dtype, device=device)

        backends = ["default"]

        if torch.device(device).type == 'cuda':
            if torch.cuda.has_magma:
                backends.append("magma")
            if has_cusolver():
                backends.append("cusolver")

        def gen_matrices():
            rhs = 3
            ns = (5, 2, 0)
            batches = ((), (0,), (1,), (2,), (2, 1), (0, 2))
            for batch, n in product(batches, ns):
                yield make_arg(batch + (n, n)), make_arg(batch + (n, rhs))
            # Shapes to exercise all the paths
            shapes = ((1, 64), (2, 128), (1025, 2))
            for b, n in shapes:
                yield make_arg((b, n, n)), make_arg((b, n, rhs))


        for A, B in gen_matrices():
            LU, pivots = torch.linalg.lu_factor(A)
            for backend in backends:
                torch.backends.cuda.preferred_linalg_library(backend)

                for left, adjoint in product((True, False), repeat=2):
                    B_left = B if left else B.mT
                    X = torch.linalg.lu_solve(LU, pivots, B_left, left=left, adjoint=adjoint)
                    A_adj = A.mH if adjoint else A
                    if left:
                        self.assertEqual(B_left, A_adj @ X)
                    else:
                        self.assertEqual(B_left, X @ A_adj)
//  [python]
def sample_inputs_lu_solve(op_info, device, dtype, requires_grad=False, **kwargs):
    """Samples the inputs for both linalg.lu_solve and lu_solve"""
    make_fn = make_fullrank_matrices_with_distinct_singular_values
    make_a = partial(make_fn, dtype=dtype, device=device)
    make_b = partial(make_tensor, dtype=dtype, device=device)

    def clone(X, requires_grad):
        Y = X.clone()
        Y.requires_grad_(requires_grad)
        return Y

    is_linalg_lu_solve = op_info.name == "linalg.lu_solve"

    batches = ((), (0,), (2,))
    ns = (3, 1, 0)
    nrhs = (4, 1, 0)

    for n, batch, rhs in product(ns, batches, nrhs):
        A = make_a(*(batch + (n, n)))
        LU, pivots = torch.lu_factor(A)

        B = make_b(batch + (n, rhs))

        grads = (False,) if not requires_grad else (True, False)
        # we try all possible combinations of requires_grad for each input
        for LU_grad, B_grad in product(grads, grads):
            # when requires_grad == True, at least one input has to have requires_grad enabled
            if requires_grad and not LU_grad and not B_grad:
                continue

            if is_linalg_lu_solve:
                for adjoint, left in product((True, False), repeat=2):
                    yield SampleInput(
                        clone(LU, LU_grad),
                        args=(pivots, clone(B if left else B.mT, B_grad)),
                        kwargs=dict(adjoint=adjoint, left=left),
                    )
            else:
                yield SampleInput(clone(B, B_grad), args=(clone(LU, LU_grad), pivots))
//  [python]
def sample_inputs_lu_solve(op_info, device, dtype, requires_grad=False, **kwargs):
    """Samples the inputs for both linalg.lu_solve and lu_solve"""
    make_fn = make_fullrank_matrices_with_distinct_singular_values
    make_a = partial(make_fn, dtype=dtype, device=device)
    make_b = partial(make_tensor, dtype=dtype, device=device)

    def clone(X, requires_grad):
        Y = X.clone()
        Y.requires_grad_(requires_grad)
        return Y

    is_linalg_lu_solve = op_info.name == "linalg.lu_solve"

    batches = ((), (0,), (2,))
    ns = (3, 1, 0)
    nrhs = (4, 1, 0)

    for n, batch, rhs in product(ns, batches, nrhs):
        A = make_a(*(batch + (n, n)))
        LU, pivots = torch.linalg.lu_factor(A)

        B = make_b(batch + (n, rhs))

        grads = (False,) if not requires_grad else (True, False)
        # we try all possible combinations of requires_grad for each input
        for LU_grad, B_grad in product(grads, grads):
            # when requires_grad == True, at least one input has to have requires_grad enabled
            if requires_grad and not LU_grad and not B_grad:
                continue

            if is_linalg_lu_solve:
                for adjoint, left in product((True, False), repeat=2):
                    yield SampleInput(
                        clone(LU, LU_grad),
                        args=(pivots, clone(B if left else B.mT, B_grad)),
                        kwargs=dict(adjoint=adjoint, left=left),
                    )
            else:
                yield SampleInput(clone(B, B_grad), args=(clone(LU, LU_grad), pivots))
//  [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_unpack(*torch.lu(A))
            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)
            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 = 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), b.lu_solve(LU_data, LU_pivots))

        sub_test(True)
        if self.device_type == 'cuda':
            sub_test(False)
//  [python]
def calculate_uncertainty_from_J(self, jacobian_matrix: torch.tensor):
        """
        Use Jacobian matrix to appraxmate Hessian matrix, then use the 
        diagnose of the inverse of Hessian matrix to calculate parameter 
        fitting uncertainty.

        Parameters
        ----------
        jacobian_matrix
            The corresponding jacobian matrix.

        Returns
        -------
        uncertainty
            The uncertainty calculated from Jacobian matrix.
        """           
        new_J, filtered_eyed_W = self.reconstruct(jacobian_matrix)
        jTj = torch.bmm(torch.bmm(new_J, filtered_eyed_W), 
                        new_J.transpose(1, 2))   
        LU, pivots = torch.linalg.lu_factor(jTj.detach())
        
        jTj_inv = torch.linalg.lu_solve(
            LU, pivots, 
            torch.eye(len(self.dataset.Profiles.variableParam))
            .to(self.dataset.device)
            .repeat(len(self.index), 1, 1)) 
        
        uncertainty = torch.pow(torch.diagonal(jTj_inv, dim1=-2, dim2=-1), 0.5)

        return uncertainty
//  [python]
def __init__(self, A=None, LU=None, P=None):
        super().__init__()
        if LU is None or P is None:
            with torch.no_grad():
                LU, P = torch.lu_unpack(*torch.lu(A))
        self.LU = LU
        self.P = P
//  [python]
def __init__(self, A=None, LU=None, P=None):
        super().__init__()
        if LU is None or P is None:
            with torch.no_grad():
                LU, P = torch.linalg.lu_factor(A)
        self.LU = LU
        self.P = P
//  [python]
def test_linalg_lu_cpu_errors(self, device, dtype):
        # Square tests
        sample = torch.randn(3, 2, 2, device=device, dtype=dtype)
        B = torch.randn(3, 2, 2, device=device, dtype=dtype)
        LU, pivots = torch.lu_factor(sample)

        # This should run without issues
        torch.linalg.lu_solve(LU, pivots, B, adjoint=True)
        torch.lu_unpack(LU, pivots)

        pivots[0] = 0
        with self.assertRaisesRegex(RuntimeError, r"greater or equal to 1"):
            torch.linalg.lu_solve(LU, pivots, B, adjoint=True)
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)

        pivots[0] = 3
        with self.assertRaisesRegex(RuntimeError, r"smaller or equal to LU.size\(-2\)"):
            torch.linalg.lu_solve(LU, pivots, B, adjoint=True)
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)

        # Rectangular tests
        sample = torch.randn(3, 4, 2, device=device, dtype=dtype)
        B = torch.randn(3, 4, 2, device=device, dtype=dtype)
        LU, pivots = torch.lu_factor(sample)

        # This should run without issues
        torch.lu_unpack(LU, pivots)

        pivots[0] = 0
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)

        pivots[0] = 5
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)


        # Rectangular tests
        sample = torch.randn(2, 3, 5, device=device, dtype=dtype)
        B = torch.randn(2, 3, 5, device=device, dtype=dtype)
        LU, pivots = torch.lu_factor(sample)

        # This should run without issues
        torch.lu_unpack(LU, pivots)

        pivots[0] = 0
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)

        pivots[0] = 4
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)
//  [python]
def test_linalg_lu_cpu_errors(self, device, dtype):
        # Square tests
        sample = torch.randn(3, 2, 2, device=device, dtype=dtype)
        B = torch.randn(3, 2, 2, device=device, dtype=dtype)
        LU, pivots = torch.linalg.lu_factor(sample)

        # This should run without issues
        torch.linalg.lu_solve(LU, pivots, B, adjoint=True)
        torch.lu_unpack(LU, pivots)

        pivots[0] = 0
        with self.assertRaisesRegex(RuntimeError, r"greater or equal to 1"):
            torch.linalg.lu_solve(LU, pivots, B, adjoint=True)
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)

        pivots[0] = 3
        with self.assertRaisesRegex(RuntimeError, r"smaller or equal to LU.size\(-2\)"):
            torch.linalg.lu_solve(LU, pivots, B, adjoint=True)
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)

        # Rectangular tests
        sample = torch.randn(3, 4, 2, device=device, dtype=dtype)
        B = torch.randn(3, 4, 2, device=device, dtype=dtype)
        LU, pivots = torch.linalg.lu_factor(sample)

        # This should run without issues
        torch.lu_unpack(LU, pivots)

        pivots[0] = 0
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)

        pivots[0] = 5
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)


        # Rectangular tests
        sample = torch.randn(2, 3, 5, device=device, dtype=dtype)
        B = torch.randn(2, 3, 5, device=device, dtype=dtype)
        LU, pivots = torch.linalg.lu_factor(sample)

        # This should run without issues
        torch.lu_unpack(LU, pivots)

        pivots[0] = 0
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)

        pivots[0] = 4
        with self.assertRaisesRegex(RuntimeError, r"between 1 and LU.size\(-2\)."):
            torch.lu_unpack(LU, pivots)
//  [python]
def _LU(M):
    """
    Perform an LU decomposition and returns L, U and a permutation vector P. 

    Args:
        M (torch.tensor): [description]

    Returns:
        tuple[torch.tensor,torch.tensor,torch.tensor]: L, U, P
    """
    LU,P = t.lu(M)
    P,L,U = t.lu_unpack(LU,P) # P transpose or not transpose?
    P = P@tn.reshape(tn.arange(P.shape[1],dtype=P.dtype,device=P.device),[-1,1])
    # P = tn.reshape(tn.arange(P.shape[1],dtype=P.dtype,device=P.device),[1,-1]) @ P
    
    return L, U, tn.squeeze(P).to(tn.int64)