jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward_no_cache(self, x, logpx=None, reverse=False, context=None):
        """Cost:
            output = O(D^2N)
            logabsdet = O(D)
        where:
            D = num of features
            N = num of inputs
        """
        lower, upper = self._create_lower_upper()
        if not reverse:
            y = F.linear(x, upper)
            y = F.linear(y, lower, self.bias)
            delta_logp = self.logabsdet() * x.new_ones(y.shape[0])
        else:
            y = x - self.bias
            y = torch.linalg.solve_triangular(lower, y.t(), upper=False, unitriangular=True)
            y = torch.linalg.solve_triangular(upper, y, upper=True, unitriangular=False)
            y = y.t()

            delta_logp = -self.logabsdet()
            delta_logp = delta_logp * x.new_ones(y.shape[0])

        if logpx is None:
            return y
        else:
            return y, logpx + delta_logp[:,None]
//  [python]
def _kl_lowrankmultivariatenormal_multivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two (Low Rank) Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (2 * q._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1) -
             _batch_lowrank_logdet(p._unbroadcasted_cov_factor, p._unbroadcasted_cov_diag,
                                   p._capacitance_tril))
    term3 = _batch_mahalanobis(q._unbroadcasted_scale_tril, (q.loc - p.loc))
    # Expands term2 according to
    # inv(qcov) @ pcov = inv(q_tril @ q_tril.T) @ (pW @ pW.T + pD)
    combined_batch_shape = torch._C._infer_size(q._unbroadcasted_scale_tril.shape[:-2],
                                                p._unbroadcasted_cov_factor.shape[:-2])
    n = p.event_shape[0]
    q_scale_tril = q._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
    p_cov_factor = p._unbroadcasted_cov_factor.expand(combined_batch_shape +
                                                      (n, p.cov_factor.size(-1)))
    p_cov_diag = (torch.diag_embed(p._unbroadcasted_cov_diag.sqrt())
                  .expand(combined_batch_shape + (n, n)))
    term21 = _batch_trace_XXT(torch.triangular_solve(p_cov_factor, q_scale_tril, upper=False)[0])
    term22 = _batch_trace_XXT(torch.triangular_solve(p_cov_diag, q_scale_tril, upper=False)[0])
    term2 = term21 + term22
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def _kl_lowrankmultivariatenormal_multivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two (Low Rank) Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (2 * q._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1) -
             _batch_lowrank_logdet(p._unbroadcasted_cov_factor, p._unbroadcasted_cov_diag,
                                   p._capacitance_tril))
    term3 = _batch_mahalanobis(q._unbroadcasted_scale_tril, (q.loc - p.loc))
    # Expands term2 according to
    # inv(qcov) @ pcov = inv(q_tril @ q_tril.T) @ (pW @ pW.T + pD)
    combined_batch_shape = torch._C._infer_size(q._unbroadcasted_scale_tril.shape[:-2],
                                                p._unbroadcasted_cov_factor.shape[:-2])
    n = p.event_shape[0]
    q_scale_tril = q._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
    p_cov_factor = p._unbroadcasted_cov_factor.expand(combined_batch_shape +
                                                      (n, p.cov_factor.size(-1)))
    p_cov_diag = (torch.diag_embed(p._unbroadcasted_cov_diag.sqrt())
                  .expand(combined_batch_shape + (n, n)))
    term21 = _batch_trace_XXT(torch.linalg.solve_triangular(p_cov_factor, q_scale_tril, upper=False))
    term22 = _batch_trace_XXT(torch.linalg.solve_triangular(p_cov_diag, q_scale_tril, upper=False))
    term2 = term21 + term22
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def _kl_lowrankmultivariatenormal_lowrankmultivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two Low Rank Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (_batch_lowrank_logdet(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                   q._capacitance_tril) -
             _batch_lowrank_logdet(p._unbroadcasted_cov_factor, p._unbroadcasted_cov_diag,
                                   p._capacitance_tril))
    term3 = _batch_lowrank_mahalanobis(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                       q.loc - p.loc,
                                       q._capacitance_tril)
    # Expands term2 according to
    # inv(qcov) @ pcov = [inv(qD) - inv(qD) @ qW @ inv(qC) @ qW.T @ inv(qD)] @ (pW @ pW.T + pD)
    #                  = [inv(qD) - A.T @ A] @ (pD + pW @ pW.T)
    qWt_qDinv = (q._unbroadcasted_cov_factor.transpose(-1, -2) /
                 q._unbroadcasted_cov_diag.unsqueeze(-2))
    A = torch.triangular_solve(qWt_qDinv, q._capacitance_tril, upper=False)[0]
    term21 = (p._unbroadcasted_cov_diag / q._unbroadcasted_cov_diag).sum(-1)
    term22 = _batch_trace_XXT(p._unbroadcasted_cov_factor *
                              q._unbroadcasted_cov_diag.rsqrt().unsqueeze(-1))
    term23 = _batch_trace_XXT(A * p._unbroadcasted_cov_diag.sqrt().unsqueeze(-2))
    term24 = _batch_trace_XXT(A.matmul(p._unbroadcasted_cov_factor))
    term2 = term21 + term22 - term23 - term24
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def _kl_lowrankmultivariatenormal_lowrankmultivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two Low Rank Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (_batch_lowrank_logdet(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                   q._capacitance_tril) -
             _batch_lowrank_logdet(p._unbroadcasted_cov_factor, p._unbroadcasted_cov_diag,
                                   p._capacitance_tril))
    term3 = _batch_lowrank_mahalanobis(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                       q.loc - p.loc,
                                       q._capacitance_tril)
    # Expands term2 according to
    # inv(qcov) @ pcov = [inv(qD) - inv(qD) @ qW @ inv(qC) @ qW.T @ inv(qD)] @ (pW @ pW.T + pD)
    #                  = [inv(qD) - A.T @ A] @ (pD + pW @ pW.T)
    qWt_qDinv = (q._unbroadcasted_cov_factor.transpose(-1, -2) /
                 q._unbroadcasted_cov_diag.unsqueeze(-2))
    A = torch.linalg.solve_triangular(q._capacitance_tril, qWt_qDinv, upper=False)
    term21 = (p._unbroadcasted_cov_diag / q._unbroadcasted_cov_diag).sum(-1)
    term22 = _batch_trace_XXT(p._unbroadcasted_cov_factor *
                              q._unbroadcasted_cov_diag.rsqrt().unsqueeze(-1))
    term23 = _batch_trace_XXT(A * p._unbroadcasted_cov_diag.sqrt().unsqueeze(-2))
    term24 = _batch_trace_XXT(A.matmul(p._unbroadcasted_cov_factor))
    term2 = term21 + term22 - term23 - term24
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def cholesky_solve(y, L):
    # L: lower triangular cholesky
    return torch.triangular_solve(
        torch.triangular_solve(y, L, upper=False)[0],
        L.conj().t(), transpose=False, upper=True
    )[0]
//  [python]
def cholesky_solve(y, L):
    # L: lower triangular cholesky
    return torch.linalg.solve_triangular(
        torch.linalg.solve_triangular(y, L, upper=False),
        L.conj().t(), upper=True
    )
//  [python]
def add_chol_rows(L11, K12, K22):
    b, m1, m2 = K12.shape
    m = m1 + m2

    L12 = torch.triangular_solve(K12, L11, upper=False)
    L22, _ = torch.cholesky_ex(K22 - (L12.mT @ L12))

    L = torch.zeros((b, m, m), device=L11.device, dtype=L11.dtype)
    L[:, :m1, :m1] = L11
    L[:, m1:, :m1] = L12
    L[:, m1:, m1:] = L22

    return L
//  [python]
def add_chol_rows(L11, K12, K22):
    b, m1, m2 = K12.shape
    m = m1 + m2

    L12 = torch.linalg.solve_triangular(K12, L11, upper=False)
    L22, _ = torch.cholesky_ex(K22 - (L12.mT @ L12))

    L = torch.zeros((b, m, m), device=L11.device, dtype=L11.dtype)
    L[:, :m1, :m1] = L11
    L[:, m1:, :m1] = L12
    L[:, m1:, m1:] = L22

    return L
//  [python]
def test_triangular_solve(self, device, dtype):
        for (k, n), (upper, unitriangular, transpose) in itertools.product(zip([2, 3, 5], [3, 5, 7]),
                                                                           itertools.product([True, False], repeat=3)):
            b, A = self.triangular_solve_test_helper((n, n), (n, k), upper,
                                                     unitriangular, device, dtype)
            x = torch.triangular_solve(b, A, upper=upper, unitriangular=unitriangular, transpose=transpose)[0]
            if transpose:
                self.assertEqual(b, A.t().mm(x))
            else:
                self.assertEqual(b, A.mm(x))
//  [python]
def test_triangular_solve(self, device, dtype):
        for (k, n), (upper, unitriangular, transpose) in itertools.product(zip([2, 3, 5], [3, 5, 7]),
                                                                           itertools.product([True, False], repeat=3)):
            b, A = self.triangular_solve_test_helper((n, n), (n, k), upper,
                                                     unitriangular, device, dtype)
            x = torch.linalg.solve_triangular(b, A, upper=upper, unitriangular=unitriangular, transpose=transpose)
            if transpose:
                self.assertEqual(b, A.t().mm(x))
            else:
                self.assertEqual(b, A.mm(x))
//  [python]
def negative_log_likelihood(self, constmean: float, log_amp: float, log_beta: torch.FloatTensor,
                                log_likelihood_noise: float) -> torch.FloatTensor:
        """
        Function used to calculate the negative log likelihood
        :param constmean:
        :param log_amp:
        :param log_beta:
        :param log_likelihood_noise:
        :return:
        """
        self.cholesky_update(constmean, log_amp, log_beta, log_likelihood_noise)

        # cholesky is lower triangular matrix
        mean_vec_sol = torch.triangular_solve(self.mean_vec, self.cholesky, upper=False)[0]
        nll = 0.5 * torch.sum(mean_vec_sol ** 2) + torch.sum(
            torch.log(torch.diag(self.cholesky))) + 0.5 * self.fit_y.size(0) * np.log(2 * np.pi)

        return nll
//  [python]
def negative_log_likelihood(self, constmean: float, log_amp: float, log_beta: torch.FloatTensor,
                                log_likelihood_noise: float) -> torch.FloatTensor:
        """
        Function used to calculate the negative log likelihood
        :param constmean:
        :param log_amp:
        :param log_beta:
        :param log_likelihood_noise:
        :return:
        """
        self.cholesky_update(constmean, log_amp, log_beta, log_likelihood_noise)

        # cholesky is lower triangular matrix
        mean_vec_sol = torch.linalg.solve_triangular(self.cholesky, self.mean_vec, upper=False)
        nll = 0.5 * torch.sum(mean_vec_sol ** 2) + torch.sum(
            torch.log(torch.diag(self.cholesky))) + 0.5 * self.fit_y.size(0) * np.log(2 * np.pi)

        return nll
//  [python]
def predict(self, x: torch.FloatTensor, **kwargs) -> (torch.FloatTensor, torch.FloatTensor):
        x = x.to(self.device, self.dtype)

        k_pred_train = self.kernel(x, self.x).evaluate()
        k_pred = self.kernel(x, diag=True)

        # cholesky is lower triangular matrix
        chol_solver = \
            torch.triangular_solve(torch.cat([k_pred_train.t(), self.mean_vec], 1), self.cholesky, upper=False)[0]
        chol_solve_k = chol_solver[:, :-1]
        chol_solve_y = chol_solver[:, -1:]

        pred_mean = torch.mm(chol_solve_k.t(), chol_solve_y) + self.mean(x).view(-1, 1)
        pred_quad = (chol_solve_k ** 2).sum(0).view(-1, 1)
        pred_var = k_pred - pred_quad
        pred_var = pred_var.clamp(min=1e-8)

        mu = self.fit_y_to_y(pred_mean).view(-1, self.num_out)
        var = (pred_var * self.y_std ** 2).view(-1, self.num_out)

        return mu, var
//  [python]
def predict(self, x: torch.FloatTensor, **kwargs) -> (torch.FloatTensor, torch.FloatTensor):
        x = x.to(self.device, self.dtype)

        k_pred_train = self.kernel(x, self.x).evaluate()
        k_pred = self.kernel(x, diag=True)

        chol_solver = \
            torch.linalg.solve_triangular(torch.cat([k_pred_train.t(), self.mean_vec], 1), self.cholesky, upper=False)
        chol_solve_k = chol_solver[:, :-1]
        chol_solve_y = chol_solver[:, -1:]

        pred_mean = torch.mm(chol_solve_k.t(), chol_solve_y) + self.mean(x).view(-1, 1)
        pred_quad = (chol_solve_k ** 2).sum(0).view(-1, 1)
        pred_var = k_pred - pred_quad
        pred_var = pred_var.clamp(min=1e-8)

        mu = self.fit_y_to_y(pred_mean).view(-1, self.num_out)
        var = (pred_var * self.y_std ** 2).view(-1, self.num_out)

        return mu, var
//  [python]
def test_block_triangular_solve(self, device, dtype):
        def run_test(a, b, upper, transpose, unitriangular, op_out):
            actual = torch.triangular_solve(b, a, upper=upper, unitriangular=unitriangular, transpose=transpose)
            actual_X = actual.solution
            actual_A_clone = actual.cloned_coefficient
            self.assertTrue(actual_A_clone.numel() == 0)
            if a._nnz() == 0:
                self.assertTrue(actual_X.isnan().all())
                return

            # TODO: replace with torch method when implemented to_dense() on block sparse tensor
            a_bsr = sp.bsr_matrix(
                (
                    a.values().cpu().numpy(),
                    a.col_indices().cpu().numpy(),
                    a.crow_indices().cpu().numpy(),
                ),
                shape=a.shape,
            )
            expected_X, _ = torch.triangular_solve(
                b,
                torch.tensor(a_bsr.todense(), device=device),
                transpose=transpose,
                upper=upper,
                unitriangular=unitriangular)
            self.assertEqual(actual_X, expected_X)

            out = torch.empty_like(b.mH if op_out and a.shape == b.shape else b)
            torch.triangular_solve(
                b, a,
                upper=upper, unitriangular=unitriangular, transpose=transpose, out=(out, actual_A_clone)
            )
            self.assertEqual(out, actual_X)
            self.assertEqual(out, expected_X)

        for index_dtype in [torch.int32, torch.int64]:
            for (m, k), block_size, noncontiguous in zip(itertools.product([1, 5], repeat=2), [2, 3], [True, False]):
                nnz = random.randint(0, m * m)
                a = self.genSparseCSRTensor((m, m), nnz, dtype=dtype, device=device, index_dtype=index_dtype)
                a_data = make_tensor((nnz, block_size, block_size), dtype=dtype, device=device)
                a_data = a_data.mT if noncontiguous else a_data  # Test column-major blocks
                a = torch._sparse_csr_tensor_unsafe(a.crow_indices(), a.col_indices(), a_data, (m * block_size, m * block_size))
                b = make_tensor((m * block_size, k), dtype=dtype, device=device, noncontiguous=noncontiguous)

                for (upper, unitriangular, transpose, op_out) in itertools.product([True, False], repeat=4):
                    run_test(a, b, upper, unitriangular, transpose, op_out)
//  [python]
def test_block_triangular_solve(self, device, dtype):
        def run_test(a, b, upper, transpose, unitriangular, op_out):
            actual = torch.linalg.solve_triangular(b, a, upper=upper, unitriangular=unitriangular, left=True, transpose=transpose)
            actual_X = actual
            actual_A_clone = torch.empty(0)
            self.assertTrue(actual_A_clone.numel() == 0)
            if a._nnz() == 0:
                self.assertTrue(actual_X.isnan().all())
                return

            # TODO: replace with torch method when implemented to_dense() on block sparse tensor
            a_bsr = sp.bsr_matrix(
                (
                    a.values().cpu().numpy(),
                    a.col_indices().cpu().numpy(),
                    a.crow_indices().cpu().numpy(),
                ),
                shape=a.shape,
            )
            expected_X = torch.linalg.solve_triangular(
                b,
                torch.tensor(a_bsr.todense(), device=device),
                upper=upper,
                left=True,
                unitriangular=unitriangular,
                transpose=transpose)
            self.assertEqual(actual_X, expected_X)

            out = torch.empty_like(b.mH if op_out and a.shape == b.shape else b)
            torch.linalg.solve_triangular(
                b, a,
                upper=upper, unitriangular=unitriangular, left=True, transpose=transpose, out=out
            )
            self.assertEqual(out, actual_X)
            self.assertEqual(out, expected_X)

        for index_dtype in [torch.int32, torch.int64]:
            for (m, k), block_size, noncontiguous in zip(itertools.product([1, 5], repeat=2), [2, 3], [True, False]):
                nnz = random.randint(0, m * m)
                a = self.genSparseCSRTensor((m, m), nnz, dtype=dtype, device=device, index_dtype=index_dtype)
                a_data = make_tensor((nnz, block_size, block_size), dtype=dtype, device=device)
                a_data = a_data.mT if noncontiguous else a_data  # Test column-major blocks
                a = torch._sparse_csr_tensor_unsafe(a.crow_indices(), a.col_indices(), a_data, (m * block_size, m * block_size))
                b = make_tensor((m * block_size, k), dtype=dtype, device=device, noncontiguous=noncontiguous)

                for (upper, unitriangular, transpose, op_out) in itertools.product([True, False], repeat=4):
                    run_test(a, b, upper, unitriangular, transpose, op_out)
//  [python]
def test_triangular_solve_batched(self, device, dtype):
        def triangular_solve_batch_helper(A_dims, b_dims, upper, unitriangular, transpose):
            b, A = self.triangular_solve_test_helper(A_dims, b_dims, upper,
                                                     unitriangular, device, dtype)
            x_exp_list = []
            for i in range(b_dims[0]):
                x_exp_list.append(torch.triangular_solve(b[i], A[i], upper=upper,
                                                         unitriangular=unitriangular,
                                                         transpose=transpose)[0])
            x_exp = torch.stack(x_exp_list)  # Stacked output
            x_act = torch.triangular_solve(b, A, upper=upper,
                                           unitriangular=unitriangular,
                                           transpose=transpose)[0]  # Actual output
            self.assertEqual(x_act, x_exp)  # Equality check
            if transpose:
                A = A.transpose(-2, -1)

            Ax = torch.matmul(A, x_act)
            self.assertEqual(b, Ax)

        for (upper, unitriangular, transpose), batchsize in itertools.product(itertools.product(
                [True, False], repeat=3), [1, 3, 4]):
            triangular_solve_batch_helper((batchsize, 5, 5), (batchsize, 5, 10),
                                          upper, unitriangular, transpose)
//  [python]
def test_triangular_solve_batched(self, device, dtype):
        def triangular_solve_batch_helper(A_dims, b_dims, upper, unitriangular, transpose):
            b, A = self.triangular_solve_test_helper(A_dims, b_dims, upper,
                                                     unitriangular, device, dtype)
            x_exp_list = []
            for i in range(b_dims[0]):
                x_exp_list.append(torch.linalg.solve_triangular(b[i], A[i], upper=upper,
                                                         unitriangular=unitriangular,
                                                         left=True,
                                                         transpose=transpose))
            x_exp = torch.stack(x_exp_list)  # Stacked output
            x_act = torch.linalg.solve_triangular(b, A, upper=upper,
                                           unitriangular=unitriangular,
                                           left=True,
                                           transpose=transpose)  # Actual output
            self.assertEqual(x_act, x_exp)  # Equality check
            if transpose:
                A = A.transpose(-2, -1)

            Ax = torch.matmul(A, x_act)
            self.assertEqual(b, Ax)

        for (upper, unitriangular, transpose), batchsize in itertools.product(itertools.product(
                [True, False], repeat=3), [1, 3, 4]):
            triangular_solve_batch_helper((batchsize, 5, 5), (batchsize, 5, 10),
                                          upper, unitriangular, transpose)
//  [python]
def test_sparse_triangular_solve(self, device, dtype):

        def run_test(n, k, upper, unitriangular, transpose, zero):
            triangle_function = torch.triu if upper else torch.tril
            make_A = torch.zeros if zero else make_tensor
            A = make_A((n, n), dtype=dtype, device=device)
            A = triangle_function(A)
            A_sparse = A.to_sparse_csr()
            B = make_tensor((n, k), dtype=dtype, device=device)

            expected = torch.triangular_solve(B, A, upper=upper, unitriangular=unitriangular, transpose=transpose)
            expected_X = expected.solution

            actual = torch.triangular_solve(B, A_sparse, upper=upper, unitriangular=unitriangular, transpose=transpose)
            actual_X = actual.solution
            actual_A_clone = actual.cloned_coefficient
            self.assertTrue(actual_A_clone.numel() == 0)
            if A_sparse._nnz() == 0:
                self.assertTrue(actual_X.isnan().all())
                return
            self.assertEqual(actual_X, expected_X)

            # test out with C contiguous strides
            out = torch.empty_strided((n, k), (k, 1), dtype=dtype, device=device)
            torch.triangular_solve(
                B, A_sparse,
                upper=upper, unitriangular=unitriangular, transpose=transpose, out=(out, actual_A_clone)
            )
            self.assertEqual(out, expected_X)

            # test out with F contiguous strides
            out = torch.empty_strided((n, k), (1, n), dtype=dtype, device=device)
            torch.triangular_solve(
                B, A_sparse,
                upper=upper, unitriangular=unitriangular, transpose=transpose, out=(out, actual_A_clone)
            )
            self.assertEqual(out, expected_X)
            self.assertEqual(out.stride(), (1, n))

            # test out with discontiguous strides
            out = torch.empty_strided((2 * n, k), (1, 2 * n), dtype=dtype, device=device)[::2]
            if n > 0 and k > 0:
                self.assertFalse(out.is_contiguous())
                self.assertFalse(out.t().is_contiguous())
            before_stride = out.stride()
            torch.triangular_solve(
                B, A_sparse,
                upper=upper, unitriangular=unitriangular, transpose=transpose, out=(out, actual_A_clone)
            )
            self.assertEqual(out, expected_X)
            self.assertEqual(out.stride(), before_stride)

        ks = [0, 1, 3]
        ns = [5, 3, 0]
        for (k, n), (upper, unitriangular, transpose, zero) in itertools.product(itertools.product(ks, ns),
                                                                                 itertools.product([True, False], repeat=4)):
            run_test(n, k, upper, unitriangular, transpose, zero)
//  [python]
def test_sparse_triangular_solve(self, device, dtype):

        def run_test(n, k, upper, unitriangular, transpose, zero):
            triangle_function = torch.triu if upper else torch.tril
            make_A = torch.zeros if zero else make_tensor
            A = make_A((n, n), dtype=dtype, device=device)
            A = triangle_function(A)
            A_sparse = A.to_sparse_csr()
            B = make_tensor((n, k), dtype=dtype, device=device)

            expected = torch.linalg.solve_triangular(B, A, upper=upper, unitriangular=unitriangular, left=True, transpose=transpose)

            actual = torch.linalg.solve_triangular(B, A_sparse, upper=upper, unitriangular=unitriangular, left=True, transpose=transpose)
            actual_X = actual
            actual_A_clone = actual.clone()
            self.assertTrue(actual_A_clone.numel() == 0)
            if A_sparse._nnz() == 0:
                self.assertTrue(actual_X.isnan().all())
                return
            self.assertEqual(actual_X, expected)

            # test out with C contiguous strides
            out = torch.empty_strided((n, k), (k, 1), dtype=dtype, device=device)
            torch.linalg.solve_triangular(
                B, A_sparse,
                upper=upper, unitriangular=unitriangular, left=True, transpose=transpose, out=out
            )
            self.assertEqual(out, expected)

            # test out with F contiguous strides
            out = torch.empty_strided((n, k), (1, n), dtype=dtype, device=device)
            torch.linalg.solve_triangular(
                B, A_sparse,
                upper=upper, unitriangular=unitriangular, left=True, transpose=transpose, out=out
            )
            self.assertEqual(out, expected)
            self.assertEqual(out.stride(), (1, n))

            # test out with discontiguous strides
            out = torch.empty_strided((2 * n, k), (1, 2 * n), dtype=dtype, device=device)[::2]
            if n > 0 and k > 0:
                self.assertFalse(out.is_contiguous())
                self.assertFalse(out.t().is_contiguous())
            before_stride = out.stride()
            torch.linalg.solve_triangular(
                B, A_sparse,
                upper=upper, unitriangular=unitriangular, left=True, transpose=transpose, out=out
            )
            self.assertEqual(out, expected)
            self.assertEqual(out.stride(), before_stride)

        ks = [0, 1, 3]
        ns = [5, 3, 0]
        for (k, n), (upper, unitriangular, transpose, zero) in itertools.product(itertools.product(ks, ns),
                                                                                 itertools.product([True, False], repeat=4)):
            run_test(n, k, upper, unitriangular, transpose, zero)
//  [python]
def test_triangular_solve_singular(self, device, dtype):
        b = torch.rand(3, 1, dtype=dtype, device=device)
        A = torch.eye(3, 3, dtype=dtype, device=device)
        A[-1, -1] = 0  # Now A is singular
        err_str = r"triangular_solve_cpu: U\(3,3\) is zero, singular U\."
        with self.assertRaisesRegex(RuntimeError, err_str):
            torch.triangular_solve(b, A)
//  [python]
def test_triangular_solve_singular(self, device, dtype):
        b = torch.rand(3, 1, dtype=dtype, device=device)
        A = torch.eye(3, 3, dtype=dtype, device=device)
        A[-1, -1] = 0  # Now A is singular
        err_str = r"triangular_solve_cpu: U\(3,3\) is zero, singular U\."
        with self.assertRaisesRegex(RuntimeError, err_str):
            torch.linalg.solve_triangular(b, A)
//  [python]
def test_triangular_solve_batched_many_batches(self, device, dtype):
        for upper, transpose, unitriangular in itertools.product([True, False], repeat=3):
            # test batched A case
            b, A = self.triangular_solve_test_helper((256, 256, 5, 5), (5, 1),
                                                     upper, unitriangular, device, dtype)
            x, _ = torch.triangular_solve(b, A,
                                          upper=upper, transpose=transpose, unitriangular=unitriangular)
            if transpose:
                A = A.transpose(-2, -1)

            Ax = torch.matmul(A, x)

            rtol = 1e-2 if dtype in [torch.float32, torch.complex64] else self.precision
            self.assertEqual(Ax, b.expand_as(Ax), atol=self.precision, rtol=rtol)

            # test batched b case
            b, A = self.triangular_solve_test_helper((3, 3), (512, 512, 3, 1),
                                                     upper, unitriangular, device, dtype)
            x, _ = torch.triangular_solve(b, A, upper=upper, transpose=transpose,
                                          unitriangular=unitriangular)
            if transpose:
                A = A.transpose(-2, -1)

            self.assertEqual(torch.matmul(A, x), b)
//  [python]
def test_triangular_solve_batched_many_batches(self, device, dtype):
        for upper, transpose, unitriangular in itertools.product([True, False], repeat=3):
            # test batched A case
            b, A = self.triangular_solve_test_helper((256, 256, 5, 5), (5, 1),
                                                     upper, unitriangular, device, dtype)
            x, _ = torch.linalg.solve_triangular(b, A,
                                          upper=upper, transpose=transpose, unitriangular=unitriangular)
            if transpose:
                A = A.transpose(-2, -1)

            Ax = torch.matmul(A, x)

            rtol = 1e-2 if dtype in [torch.float32, torch.complex64] else self.precision
            self.assertEqual(Ax, b.expand_as(Ax), atol=self.precision, rtol=rtol)

            # test batched b case
            b, A = self.triangular_solve_test_helper((3, 3), (512, 512, 3, 1),
                                                     upper, unitriangular, device, dtype)
            x, _ = torch.linalg.solve_triangular(b, A, upper=upper, transpose=transpose,
                                          unitriangular=unitriangular)
            if transpose:
                A = A.transpose(-2, -1)

            self.assertEqual(torch.matmul(A, x), b)
//  [python]
def test_triangular_solve_autograd(self, device, dtype):
        def run_test(A_dims, B_dims):
            A = torch.rand(*A_dims, dtype=dtype).requires_grad_()
            b = torch.rand(*B_dims, dtype=dtype).requires_grad_()

            for upper, transpose, unitriangular in itertools.product((True, False), repeat=3):
                def func(A, b):
                    return torch.triangular_solve(b, A, upper, transpose, unitriangular)

                gradcheck(func, [A, b])
                gradgradcheck(func, [A, b])

        run_test((3, 3), (3, 4))
        run_test((3, 3), (3, 2))
        run_test((2, 3, 3), (2, 3, 4))
        run_test((2, 3, 3), (2, 3, 2))
//  [python]
def test_triangular_solve_autograd(self, device, dtype):
        def run_test(A_dims, B_dims):
            A = torch.rand(*A_dims, dtype=dtype).requires_grad_()
            b = torch.rand(*B_dims, dtype=dtype).requires_grad_()

            for upper, transpose, unitriangular in itertools.product((True, False), repeat=3):
                def func(A, b):
                    return torch.linalg.solve_triangular(b, A, upper=upper, transpose=transpose, unitriangular=unitriangular)

                gradcheck(func, [A, b])
                gradgradcheck(func, [A, b])

        run_test((3, 3), (3, 4))
        run_test((3, 3), (3, 2))
        run_test((2, 3, 3), (2, 3, 4))
        run_test((2, 3, 3), (2, 3, 2))
//  [python]
def inverse(self, inputs):
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        outputs, _ = torch.triangular_solve(outputs.t(), lower, upper=False, unitriangular=False)
        outputs = outputs.t()
        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])
        return outputs, logabsdet
//  [python]
def inverse(self, inputs):
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        outputs, _ = torch.linalg.solve_triangular(outputs.t(), lower, upper=False, unitriangular=False)
        outputs = outputs.t()
        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])
        return outputs, logabsdet
//  [python]
def _inverse(self, y):
        """
        Inverts y => x.
        :param y: (*sample_shape, input_dim) output of the bijection
        :type y: torch.Tensor
        """
        # PLUx = y
        # LUx = P^T y
        y = y[..., None]  # (*sample_shape, input_dim, 1)
        LUx = torch.matmul(self.permutation.T, y)  # (*sample_shape, input_dim, 1)

        # Solve L(Ux) = P^T y
        Ux, _ = torch.triangular_solve(LUx, self.L, upper=False)  # (*sample_shape, input_dim, 1)

        # Solve Ux = (PL)^-1y
        x, _ = torch.triangular_solve(Ux, self.U)  # (*sample_shape, input_dim, 1)
        # print('\nIN INVERSE')
        # print('y max min', y[..., 0].max().item(), y[..., 0].min().item())
        # print('x max min', x[..., 0].max().item(), x[..., 0].min().item())
        # print('determinant', (self.permutation @ self.L @ self.U).det())
        # print('determinant inverse', (self.permutation @ self.L @ self.U).inverse().det())
        # print('2-norm', (self.permutation @ self.L @ self.U).norm(dim=(-1, -2), p=2))
        # print('2-norm inverse', (self.permutation @ self.L @ self.U).inverse().norm(p=2, dim=(-1, -2)))
        return x[..., 0]
//  [python]
def _inverse(self, y):
        """
        Inverts y => x.
        :param y: (*sample_shape, input_dim) output of the bijection
        :type y: torch.Tensor
        """
        # PLUx = y
        # LUx = P^T y
        y = y[..., None]  # (*sample_shape, input_dim, 1)
        LUx = torch.matmul(self.permutation.T, y)  # (*sample_shape, input_dim, 1)

        # Solve L(Ux) = P^T y
        Ux = torch.linalg.solve_triangular(self.L, LUx, upper=False)  # (*sample_shape, input_dim, 1)

        # Solve Ux = (PL)^-1y
        x = torch.linalg.solve_triangular(self.U, Ux)  # (*sample_shape, input_dim, 1)
        return x[..., 0]
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Performs whitening using the Cholesky decomposition.

        Args:
            x (torch.Tensor): a batch or slice of projected features.

        Returns:
            torch.Tensor: a batch or slice of whitened features.
        """

        x = x.unsqueeze(2).unsqueeze(3)
        m = x.mean(0).view(self.output_dim, -1).mean(-1).view(1, -1, 1, 1)
        xn = x - m

        T = xn.permute(1, 0, 2, 3).contiguous().view(self.output_dim, -1)
        f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)

        eye = torch.eye(self.output_dim).type(f_cov.type())

        f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye

        inv_sqrt = torch.triangular_solve(eye, torch.linalg.cholesky(f_cov_shrinked), upper=False)[
            0
        ]
        inv_sqrt = inv_sqrt.contiguous().view(self.output_dim, self.output_dim, 1, 1)

        decorrelated = conv2d(xn, inv_sqrt)

        return decorrelated.squeeze(2).squeeze(2)
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Performs whitening using the Cholesky decomposition.

        Args:
            x (torch.Tensor): a batch or slice of projected features.

        Returns:
            torch.Tensor: a batch or slice of whitened features.
        """

        x = x.unsqueeze(2).unsqueeze(3)
        m = x.mean(0).view(self.output_dim, -1).mean(-1).view(1, -1, 1, 1)
        xn = x - m

        T = xn.permute(1, 0, 2, 3).contiguous().view(self.output_dim, -1)
        f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)

        eye = torch.eye(self.output_dim).type(f_cov.type())

        f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye

        inv_sqrt = torch.linalg.solve_triangular(eye, torch.linalg.cholesky(f_cov_shrinked), upper=False)
        inv_sqrt = inv_sqrt.contiguous().view(self.output_dim, self.output_dim, 1, 1)

        decorrelated = conv2d(xn, inv_sqrt)

        return decorrelated.squeeze(2).squeeze(2)
//  [python]
def test_triangular_solve(self, device, dtype):
        ks = [0, 1, 3]
        ns = [0, 5]
        for (k, n), (upper, unitriangular, transpose) in itertools.product(zip(ks, ns),
                                                                           itertools.product([True, False], repeat=3)):
            b, A = self.triangular_solve_test_helper((n, n), (n, k), upper,
                                                     unitriangular, device, dtype)
            x = torch.triangular_solve(b, A, upper=upper, unitriangular=unitriangular, transpose=transpose)[0]
            if transpose:
                self.assertEqual(b, np.matmul(A.t().cpu(), x.cpu()))
            else:
                self.assertEqual(b, np.matmul(A.cpu(), x.cpu()))
//  [python]
def test_triangular_solve(self, device, dtype):
        ks = [0, 1, 3]
        ns = [0, 5]
        for (k, n), (upper, unitriangular, transpose) in itertools.product(zip(ks, ns),
                                                                           itertools.product([True, False], repeat=3)):
            b, A = self.triangular_solve_test_helper((n, n), (n, k), upper,
                                                     unitriangular, device, dtype)
            x = torch.linalg.solve_triangular(b, A, upper=upper, unitriangular=unitriangular, transpose=transpose)
            if transpose:
                self.assertEqual(b, np.matmul(A.t().cpu(), x.cpu()))
            else:
                self.assertEqual(b, np.matmul(A.cpu(), x.cpu()))
//  [python]
def test_triangular_solve_singular(self, device, dtype):
        b = torch.rand(3, 1, dtype=dtype, device=device)
        A = torch.eye(3, 3, dtype=dtype, device=device)
        A[-1, -1] = 0  # Now A is singular
        err_str = r"triangular_solve: U\(3,3\) is zero, singular U\."
        with self.assertRaisesRegex(RuntimeError, err_str):
            torch.triangular_solve(b, A)
//  [python]
def test_triangular_solve_singular(self, device, dtype):
        b = torch.rand(3, 1, dtype=dtype, device=device)
        A = torch.eye(3, 3, dtype=dtype, device=device)
        A[-1, -1] = 0  # Now A is singular
        err_str = r"triangular_solve: U\(3,3\) is zero, singular U\."
        with self.assertRaisesRegex(RuntimeError, err_str):
            torch.linalg.solve_triangular(b, A)
//  [python]
def test_triangular_solve_batched_broadcasting(self, device, dtype):
        from scipy.linalg import solve_triangular as tri_solve

        def scipy_tri_solve_batched(A, B, upper, trans, diag):
            batch_dims_A, batch_dims_B = A.shape[:-2], B.shape[:-2]
            single_dim_A, single_dim_B = A.shape[-2:], B.shape[-2:]
            expand_dims = tuple(torch._C._infer_size(torch.Size(batch_dims_A),
                                                     torch.Size(batch_dims_B)))
            expand_A = np.broadcast_to(A, expand_dims + single_dim_A)
            expand_B = np.broadcast_to(B, expand_dims + single_dim_B)
            flat_A = expand_A.reshape((-1,) + single_dim_A)
            flat_B = expand_B.reshape((-1,) + single_dim_B)
            flat_X = np.vstack([tri_solve(a, b, lower=(not upper), trans=int(trans), unit_diagonal=diag)
                                for a, b in zip(flat_A, flat_B)])
            return flat_X.reshape(expand_B.shape)

        def run_test(A_dims, b_dims, device, upper, transpose, unitriangular):
            b, A = self.triangular_solve_test_helper(A_dims, b_dims, upper,
                                                     unitriangular, device, dtype)
            x_exp = torch.as_tensor(scipy_tri_solve_batched(A.cpu().numpy(), b.cpu().numpy(),
                                                            upper, transpose, unitriangular))
            x = torch.triangular_solve(b, A, upper=upper, transpose=transpose, unitriangular=unitriangular)[0]

            self.assertEqual(x, x_exp.to(device))

        for upper, transpose, unitriangular in itertools.product([True, False], repeat=3):
            # test against scipy.linalg.solve_triangular
            run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6), device, upper, transpose, unitriangular)  # no broadcasting
            run_test((2, 1, 3, 4, 4), (4, 6), device, upper, transpose, unitriangular)  # broadcasting b
            run_test((4, 4), (2, 1, 3, 4, 2), device, upper, transpose, unitriangular)  # broadcasting A
            run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5), device, upper, transpose, unitriangular)
//  [python]
def test_triangular_solve_batched_broadcasting(self, device, dtype):
        from scipy.linalg import solve_triangular as tri_solve

        def scipy_tri_solve_batched(A, B, upper, trans, diag):
            batch_dims_A, batch_dims_B = A.shape[:-2], B.shape[:-2]
            single_dim_A, single_dim_B = A.shape[-2:], B.shape[-2:]
            expand_dims = tuple(torch._C._infer_size(torch.Size(batch_dims_A),
                                                     torch.Size(batch_dims_B)))
            expand_A = np.broadcast_to(A, expand_dims + single_dim_A)
            expand_B = np.broadcast_to(B, expand_dims + single_dim_B)
            flat_A = expand_A.reshape((-1,) + single_dim_A)
            flat_B = expand_B.reshape((-1,) + single_dim_B)
            flat_X = np.vstack([tri_solve(a, b, lower=(not upper), trans=int(trans), unit_diagonal=diag)
                                for a, b in zip(flat_A, flat_B)])
            return flat_X.reshape(expand_B.shape)

        def run_test(A_dims, b_dims, device, upper, transpose, unitriangular):
            b, A = self.triangular_solve_test_helper(A_dims, b_dims, upper,
                                                     unitriangular, device, dtype)
            x_exp = torch.as_tensor(scipy_tri_solve_batched(A.cpu().numpy(), b.cpu().numpy(),
                                                            upper, transpose, unitriangular))
            x = torch.linalg.solve_triangular(b, A, upper=upper, transpose=transpose, unitriangular=unitriangular)

            self.assertEqual(x, x_exp.to(device))

        for upper, transpose, unitriangular in itertools.product([True, False], repeat=3):
            # test against scipy.linalg.solve_triangular
            run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6), device, upper, transpose, unitriangular)  # no broadcasting
            run_test((2, 1, 3, 4, 4), (4, 6), device, upper, transpose, unitriangular)  # broadcasting b
            run_test((4, 4), (2, 1, 3, 4, 2), device, upper, transpose, unitriangular)  # broadcasting A
            run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5), device, upper, transpose, unitriangular)
//  [python]
def predict(self, xstar, y, x, full_cov=False):
        """
        compute posterior p(f* | x, y)
        """

        prec = self._dist(x.squeeze()).precision_matrix  #(K+sigma^2I)^-1
        m = x.shape[-1]
        d = x.shape[-2]

        x = self.scale * self.dim_scale * x
        xstar = self.scale * self.dim_scale * xstar

        x = x[..., None, :, :]  #(...,d,m)
        xstar = xstar[..., None, :, :]  #(...,d,m)
        xt = x.transpose(-1, -2)  #(...,m,d)
        variance = self.prms  #(n) (p(Y|F) variance)
        cov_diag = variance[..., None] * torch.ones(m)  #(n x m)
        capacitance_tril = batch_capacitance_tril(xt, cov_diag)

        xdinv = (x / variance[..., None, None])
        A = torch.triangular_solve(xdinv, capacitance_tril,
                                   upper=False)[0]  #(...,d,m)
        y = y[..., None]
        _mu1 = xdinv.matmul(y)
        _mu2 = x.matmul(A.transpose(-1, -2).matmul(A.matmul(y)))
        mu = xstar.transpose(-1, -2).matmul(_mu1 - _mu2).squeeze(-1)
        if not full_cov:
            v1 = torch.square(xstar).sum(-2)
            v2 = xstar.transpose(-1, -2).matmul(
                x / variance[:, None, None].sqrt()).square().sum(-1)
            v3 = A.matmul(x.transpose(-1, -2)).matmul(xstar).square().sum(-2)
            v = v1 - v2 + v3
            return mu, v1 - v2 + v3
        else:
            xxT = xdinv.matmul(x.transpose(-1, -2))
            xAT = x.matmul(A.transpose(-1, -2))
            xATAxT = xAT.matmul(xAT.transpose(-1, -2))
            z = torch.eye(d) - xxT + xATAxT
            c = xstar.transpose(-1, -2).matmul(z).matmul(xstar)
            return mu, c
//  [python]
def predict(self, xstar, y, x, full_cov=False):
        """
        compute posterior p(f* | x, y)
        """

        prec = self._dist(x.squeeze()).precision_matrix  #(K+sigma^2I)^-1
        m = x.shape[-1]
        d = x.shape[-2]

        x = self.scale * self.dim_scale * x
        xstar = self.scale * self.dim_scale * xstar

        x = x[..., None, :, :]  #(...,d,m)
        xstar = xstar[..., None, :, :]  #(...,d,m)
        xt = x.transpose(-1, -2)  #(...,m,d)
        variance = self.prms  #(n) (p(Y|F) variance)
        cov_diag = variance[..., None] * torch.ones(m)  #(n x m)
        capacitance_tril = batch_capacitance_tril(xt, cov_diag)

        xdinv = (x / variance[..., None, None])
        A = torch.linalg.solve_triangular(xdinv, capacitance_tril,
                                   upper=False)  #(...,d,m)
        y = y[..., None]
        _mu1 = xdinv.matmul(y)
        _mu2 = x.matmul(A.transpose(-1, -2).matmul(A.matmul(y)))
        mu = xstar.transpose(-1, -2).matmul(_mu1 - _mu2).squeeze(-1)
        if not full_cov:
            v1 = torch.square(xstar).sum(-2)
            v2 = xstar.transpose(-1, -2).matmul(
                x / variance[:, None, None].sqrt()).square().sum(-1)
            v3 = A.matmul(x.transpose(-1, -2)).matmul(xstar).square().sum(-2)
            v = v1 - v2 + v3
            return mu, v1 - v2 + v3
        else:
            xxT = xdinv.matmul(x.transpose(-1, -2))
            xAT = x.matmul(A.transpose(-1, -2))
            xATAxT = xAT.matmul(xAT.transpose(-1, -2))
            z = torch.eye(d) - xxT + xATAxT
            c = xstar.transpose(-1, -2).matmul(z).matmul(xstar)
            return mu, c
//  [python]
def precompute_backward(self, delta):
        return torch.triangular_solve(self.I, self.I - delta*self.A, upper=False)[0]
//  [python]
def precompute_backward(self, delta):
        return torch.linalg.solve_triangular(self.I, self.I - delta*self.A, upper=False)
//  [python]
def memoized_cholesky_decomposition(self):
        "memoize the cholesky decomposition to avoid recomputing it when we are not training or when we are fitting it"
        if (self._L_train_train is None) or (self.training):
            # covariance between training samples
            cov_train_train = self.kernel.matrix((self.train_input_cat, self.train_input_cont), (self.train_input_cat, self.train_input_cont))
            # cholesky decompositions (accelerate solving of linear systems)
            self._L_train_train = psd_safe_cholesky(cov_train_train).detach() # we drop the gradient for the cholesky decomposition
            # outputs for the training data with prior correction
            train_outputs = self.train_outputs - self.prior(self.train_input_cat, self.train_input_cont)
            # weights for the predicted mean
            self._output_weights, _ = torch.triangular_solve(train_outputs, self._L_train_train, upper=False)
        return self._L_train_train, self._output_weights
//  [python]
def memoized_cholesky_decomposition(self):
        "memoize the cholesky decomposition to avoid recomputing it when we are not training or when we are fitting it"
        if (self._L_train_train is None) or (self.training):
            # covariance between training samples
            cov_train_train = self.kernel.matrix((self.train_input_cat, self.train_input_cont), (self.train_input_cat, self.train_input_cont))
            # cholesky decompositions (accelerate solving of linear systems)
            self._L_train_train = psd_safe_cholesky(cov_train_train).detach() # we drop the gradient for the cholesky decomposition
            # outputs for the training data with prior correction
            train_outputs = self.train_outputs - self.prior(self.train_input_cat, self.train_input_cont)
            # weights for the predicted mean
            self._output_weights, _ = torch.linalg.solve_triangular(train_outputs, self._L_train_train, upper=False)
        return self._L_train_train, self._output_weights