jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _triangular_inv(x, upper=False):
    if x.size(-1) == 1:
        return x.reciprocal()
    return torch.triangular_solve(torch.eye(x.size(-1)), x, upper=upper).solution
//  [python]
def _triangular_inv(x, upper=False):
    if x.size(-1) == 1:
        return x.reciprocal()
    return torch.linalg.solve_triangular(x, torch.eye(x.size(-1)), upper=upper)
//  [python]
def test_solver_gradient_coo_triu(self):
        # Sparse solver:
        As1 = self.As_coo_triu.detach().clone()
        As1.requires_grad = True
        Bd1 = self.Bd.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = self.solve(As1, Bd1, upper=True, unitriangular=self.unitriangular)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.Ad_triu.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.Bd.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.solve(Bd2, Ad2, upper=True).solution
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_solver_gradient_coo_triu(self):
        # Sparse solver:
        As1 = self.As_coo_triu.detach().clone()
        As1.requires_grad = True
        Bd1 = self.Bd.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = self.solve(As1, Bd1, upper=True, unitriangular=self.unitriangular)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.Ad_triu.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.Bd.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.linalg.solve_triangular(Ad2, Bd2, upper=True, unitriangular=self.unitriangular)
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def precision_to_scale_tril(P: tc.Tensor) -> tc.Tensor:
    # Ref: https://nbviewer.jupyter.org/gist/fehiepsi/5ef8e09e61604f10607380467eb82006#Precision-to-scale_tril
    Lf = tc.cholesky(tc.flip(P, dims=(-2, -1)))
    L_inv = tc.transpose(tc.flip(Lf, dims=(-2, -1)), -2, -1)
    L = tc.triangular_solve(tc.eye(P.shape[-1], dtype=P.dtype, device=P.device),
            L_inv, upper=False)[0]
    return L
//  [python]
def precision_to_scale_tril(P: tc.Tensor) -> tc.Tensor:
    # Ref: https://nbviewer.jupyter.org/gist/fehiepsi/5ef8e09e61604f10607380467eb82006#Precision-to-scale_tril
    Lf = tc.linalg.cholesky(tc.flip(P, dims=(-2, -1)))
    L_inv = tc.transpose(tc.flip(Lf, dims=(-2, -1)), -2, -1)
    L = tc.triangular_solve(tc.eye(P.shape[-1], dtype=P.dtype, device=P.device),
            L_inv, upper=False)[0]
    return L
//  [python]
def test_solver_gradient_coo_tril(self):
        # Sparse solver:
        As1 = self.As_coo_tril.detach().clone()
        As1.requires_grad = True
        Bd1 = self.Bd.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = self.solve(As1, Bd1, upper=False, unitriangular=self.unitriangular)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.Ad_tril.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.Bd.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.solve(Bd2, Ad2).solution
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_solver_gradient_coo_tril(self):
        # Sparse solver:
        As1 = self.As_coo_tril.detach().clone()
        As1.requires_grad = True
        Bd1 = self.Bd.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = self.solve(As1, Bd1, upper=False, unitriangular=self.unitriangular)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.Ad_tril.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.Bd.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.linalg.solve_triangular(Ad2, Bd2, upper=False, unitriangular=self.unitriangular)
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_solver_gradient_csr_triu(self):
        # Sparse solver:
        As1 = self.As_csr_triu.detach().clone()
        As1.requires_grad = True
        Bd1 = self.Bd.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = self.solve(As1, Bd1, upper=True, unitriangular=self.unitriangular)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.Ad_triu.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.Bd.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.solve(Bd2, Ad2, upper=True).solution
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_solver_gradient_csr_triu(self):
        # Sparse solver:
        As1 = self.As_csr_triu.detach().clone()
        As1.requires_grad = True
        Bd1 = self.Bd.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = self.solve(As1, Bd1, upper=True, unitriangular=self.unitriangular)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.Ad_triu.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.Bd.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.linalg.solve_triangular(Ad2, Bd2, upper=True, unitriangular=self.unitriangular)
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def _batch_capacitance_tril(W, D):
    r"""
    Computes Cholesky of :math:`I + W.T @ inv(D) @ W` for a batch of matrices :math:`W`
    and a batch of vectors :math:`D`.
    """
    m = W.size(-1)
    Wt_Dinv = W.transpose(-1, -2) / D.unsqueeze(-2)
    K = torch.matmul(Wt_Dinv, W).contiguous()
    K.view(-1, m * m)[:, ::m + 1] += 1  # add identity matrix to K
    return torch.cholesky(K)
//  [python]
def _batch_capacitance_tril(W, D):
    r"""
    Computes Cholesky of :math:`I + W.T @ inv(D) @ W` for a batch of matrices :math:`W`
    and a batch of vectors :math:`D`.
    """
    m = W.size(-1)
    Wt_Dinv = W.transpose(-1, -2) / D.unsqueeze(-2)
    K = torch.matmul(Wt_Dinv, W).contiguous()
    K.view(-1, m * m)[:, ::m + 1] += 1  # add identity matrix to K
    return torch.linalg.cholesky(K)
//  [python]
def forward(self, x, numpy=False):
        if self.training:
            mean = x.mean(dim=0)
            x = x - mean
            cov = x.t().matmul(x) / (x.size(0) - 1)
            self.running_mean = self.momentum * mean + (1 - self.momentum) * self.running_mean
            self.running_covariance = self.momentum * cov + (1 - self.momentum) * self.running_covariance
        else:
            mean = self.running_mean
            cov = self.running_covariance
            x = x - mean

        cov = (1 - self.eps) * cov + self.eps * torch.eye(self.num_features).to(cov)
        if numpy:

            I = torch.eye(x.size(1)).to(cov).detach().cpu().numpy()
            cv = np.linalg.cholesky(cov.detach().cpu().numpy())
            whitening_transform = solve_triangular(cv, I, lower=True).T
            
            whitening_transform = torch.tensor(whitening_transform).to(x)
        else:
            I = torch.eye(x.size(1)).to(cov).cpu()
            C = torch.cholesky(cov.cpu())
            whitening_transform = torch.triangular_solve(I, C, upper=False)[0].t().to(x.device)
        return x.matmul(whitening_transform)
//  [python]
def forward(self, x, numpy=False):
        if self.training:
            mean = x.mean(dim=0)
            x = x - mean
            cov = x.t().matmul(x) / (x.size(0) - 1)
            self.running_mean = self.momentum * mean + (1 - self.momentum) * self.running_mean
            self.running_covariance = self.momentum * cov + (1 - self.momentum) * self.running_covariance
        else:
            mean = self.running_mean
            cov = self.running_covariance
            x = x - mean

        cov = (1 - self.eps) * cov + self.eps * torch.eye(self.num_features).to(cov)
        if numpy:

            I = torch.eye(x.size(1)).to(cov).detach().cpu().numpy()
            cv = np.linalg.cholesky(cov.detach().cpu().numpy())
            whitening_transform = solve_triangular(cv, I, lower=True).T
            
            whitening_transform = torch.tensor(whitening_transform).to(x)
        else:
            I = torch.eye(x.size(1)).to(cov).cpu()
            C = torch.linalg.cholesky(cov.cpu())
            whitening_transform = torch.triangular_solve(I, C, upper=False)[0].t().to(x.device)
        return x.matmul(whitening_transform)
//  [python]
def test_solver_gradient_csr_tril(self):
        # Sparse solver:
        As1 = self.As_csr_tril.detach().clone()
        As1.requires_grad = True
        Bd1 = self.Bd.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = self.solve(As1, Bd1, upper=False, unitriangular=self.unitriangular)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.Ad_tril.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.Bd.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.solve(Bd2, Ad2).solution
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_solver_gradient_csr_tril(self):
        # Sparse solver:
        As1 = self.As_csr_tril.detach().clone()
        As1.requires_grad = True
        Bd1 = self.Bd.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = self.solve(As1, Bd1, upper=False, unitriangular=self.unitriangular)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.Ad_tril.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.Bd.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.linalg.solve_triangular(Ad2, Bd2, upper=False, unitriangular=self.unitriangular)
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def batch_cholesky_inverse(matrix):
	"""
	Args: 	matrix, 4-D tensor, [N, C, M, M].
			matrix must be a symmetric positive define matrix.
	"""
	chol_low = torch.cholesky(matrix, upper=False)
	chol_low_inv = batch_low_tri_inv(chol_low)
	return torch.matmul(chol_low_inv.transpose(-2, -1), chol_low_inv)
//  [python]
def batch_cholesky_inverse(matrix):
	"""
	Args: 	matrix, 4-D tensor, [N, C, M, M].
			matrix must be a symmetric positive define matrix.
	"""
	chol_low = torch.linalg.cholesky(matrix, upper=False)
	chol_low_inv = batch_low_tri_inv(chol_low)
	return torch.matmul(chol_low_inv.transpose(-2, -1), chol_low_inv)
//  [python]
def cholesky(self):
        cov = self.covariance_matrix
        return torch.cholesky(cov)
//  [python]
def cholesky(self):
        cov = self.covariance_matrix
        return torch.linalg.cholesky(cov)
//  [python]
def log_det_by_cholesky(matrix):
	"""
	Args:
		matrix: matrix must be a positive define matrix.
				shape [N, C, D, D].
	Ref:
		https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
	"""
	# This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
	# where C is the cholesky decomposition of A.
	chol = torch.cholesky(matrix)
	#return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-6), dim=-1)
	return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def log_det_by_cholesky(matrix):
	"""
	Args:
		matrix: matrix must be a positive define matrix.
				shape [N, C, D, D].
	Ref:
		https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
	"""
	# This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
	# where C is the cholesky decomposition of A.
	chol = torch.linalg.cholesky(matrix)
	#return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-6), dim=-1)
	return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def scale_tril(self):
        # The following identity is used to increase the numerically computation stability
        # for Cholesky decomposition (see http://www.gaussianprocess.org/gpml/, Section 3.4.3):
        #     W @ W.T + D = D1/2 @ (I + D-1/2 @ W @ W.T @ D-1/2) @ D1/2
        # The matrix "I + D-1/2 @ W @ W.T @ D-1/2" has eigenvalues bounded from below by 1,
        # hence it is well-conditioned and safe to take Cholesky decomposition.
        n = self._event_shape[0]
        cov_diag_sqrt_unsqueeze = self._unbroadcasted_cov_diag.sqrt().unsqueeze(-1)
        Dinvsqrt_W = self._unbroadcasted_cov_factor / cov_diag_sqrt_unsqueeze
        K = torch.matmul(Dinvsqrt_W, Dinvsqrt_W.transpose(-1, -2)).contiguous()
        K.view(-1, n * n)[:, ::n + 1] += 1  # add identity matrix to K
        scale_tril = cov_diag_sqrt_unsqueeze * torch.cholesky(K)
        return scale_tril.expand(self._batch_shape + self._event_shape + self._event_shape)
//  [python]
def scale_tril(self):
        # The following identity is used to increase the numerically computation stability
        # for Cholesky decomposition (see http://www.gaussianprocess.org/gpml/, Section 3.4.3):
        #     W @ W.T + D = D1/2 @ (I + D-1/2 @ W @ W.T @ D-1/2) @ D1/2
        # The matrix "I + D-1/2 @ W @ W.T @ D-1/2" has eigenvalues bounded from below by 1,
        # hence it is well-conditioned and safe to take Cholesky decomposition.
        n = self._event_shape[0]
        cov_diag_sqrt_unsqueeze = self._unbroadcasted_cov_diag.sqrt().unsqueeze(-1)
        Dinvsqrt_W = self._unbroadcasted_cov_factor / cov_diag_sqrt_unsqueeze
        K = torch.matmul(Dinvsqrt_W, Dinvsqrt_W.transpose(-1, -2)).contiguous()
        K.view(-1, n * n)[:, ::n + 1] += 1  # add identity matrix to K
        scale_tril = cov_diag_sqrt_unsqueeze * torch.linalg.cholesky(K)
        return scale_tril.expand(self._batch_shape + self._event_shape + self._event_shape)
//  [python]
def forward(ctx, H, b):
        # don't crash training if cholesky decomp fails
        try:
            U = torch.cholesky(H)
            xs = torch.cholesky_solve(b, U)
            ctx.save_for_backward(U, xs)
            ctx.failed = False
        except Exception as e:
            print(e)
            ctx.failed = True
            xs = torch.zeros_like(b)

        return xs
//  [python]
def forward(ctx, H, b):
        # don't crash training if cholesky decomp fails
        try:
            U = torch.linalg.cholesky(H)
            xs = torch.linalg.solve_triangular(U, b)
            ctx.save_for_backward(U, xs)
            ctx.failed = False
        except Exception as e:
            print(e)
            ctx.failed = True
            xs = torch.zeros_like(b)

        return xs
//  [python]
def cholesky(x, *args, **kwargs):
    r"""A wrapper over `torch.cholesky` that deals with the CUDA maximum batch
    size limitation.
    """
    if x.device.type == 'cpu':
        return torch.cholesky(x, *args, **kwargs)

    shape = x.shape
    batch_size = np.prod(shape[:-2])
    cuda_batch_size = (1 << 18) - 4

    if batch_size <= cuda_batch_size:
        return torch.cholesky(x, *args, **kwargs)

    x = x.view(-1, shape[-1], shape[-1])
    acc = torch.empty_like(x)
    for i in range(0, batch_size, cuda_batch_size):
        indices = torch.arange(
                i,
                min(batch_size, i + cuda_batch_size),
                device=x.device,
                dtype=torch.int64)
        acc[indices] = torch.cholesky(x[indices], *args, **kwargs)

    return acc.view(shape)
//  [python]
def cholesky(x, *args, **kwargs):
    r"""A wrapper over `torch.linalg.cholesky` that deals with the CUDA maximum batch
    size limitation.
    """
    if x.device.type == 'cpu':
        return torch.linalg.cholesky(x, *args, **kwargs)

    shape = x.shape
    batch_size = np.prod(shape[:-2])
    cuda_batch_size = (1 << 18) - 4

    if batch_size <= cuda_batch_size:
        return torch.linalg.cholesky(x, *args, **kwargs)

    x = x.view(-1, shape[-1], shape[-1])
    acc = torch.empty_like(x)
    for i in range(0, batch_size, cuda_batch_size):
        indices = torch.arange(
                i,
                min(batch_size, i + cuda_batch_size),
                device=x.device,
                dtype=torch.int64)
        acc[indices] = torch.linalg.cholesky(x[indices], *args, **kwargs)

    return acc.view(shape)
//  [python]
def sample_MVG(M, U, v):
    # U is a matrix, v is the diagonal of V
    # M: rxc
    # U: rxr
    # V: cxc

    A = torch.cholesky(U + 1e-3*torch.diag(torch.ones(U.shape[0])))
    #A = torch.cholesky(U)
    B = torch.diag(v.sqrt())
    E = torch.randn(M.shape)
    
    return M + torch.mm(torch.mm(A, E), B)
//  [python]
def sample_MVG(M, U, v):
    # U is a matrix, v is the diagonal of V
    # M: rxc
    # U: rxr
    # V: cxc

    A = torch.linalg.cholesky(U + 1e-3*torch.diag(torch.ones(U.shape[0])))
    #A = torch.linalg.cholesky(U)
    B = torch.diag(v.sqrt())
    E = torch.randn(M.shape)
    
    return M + torch.mm(torch.mm(A, E), B)
//  [python]
def potrf(a, upper=True, out=None):
    r"""Computes the Cholesky decomposition of a symmetric positive-definite
    matrix :math:`A`.

    For more information, regarding :func:`torch.potrf`, please check :func:`torch.cholesky`.

    .. warning::
        torch.potrf is deprecated in favour of torch.cholesky and will be removed in the next
        release. Please use torch.cholesky instead and note that the :attr:`upper` argument in
        torch.cholesky defaults to ``False``.
    """
    warnings.warn("torch.potrf is deprecated in favour of torch.cholesky and will be removed in the next "
                  "release. Please use torch.cholesky instead and note that the :attr:`upper` argument in"
                  " torch.cholesky defaults to ``False``.", stacklevel=2)
    return torch.cholesky(a, upper=upper, out=out)
//  [python]
def potrf(a, upper=True, out=None):
    r"""Computes the Cholesky decomposition of a symmetric positive-definite
    matrix :math:`A`.

    For more information, regarding :func:`torch.potrf`, please check :func:`torch.cholesky`.

    .. warning::
        torch.potrf is deprecated in favour of torch.cholesky and will be removed in the next
        release. Please use torch.cholesky instead and note that the :attr:`upper` argument in
        torch.cholesky defaults to ``False``.
    """
    warnings.warn("torch.potrf is deprecated in favour of torch.cholesky and will be removed in the next "
                  "release. Please use torch.cholesky instead and note that the :attr:`upper` argument in"
                  " torch.cholesky defaults to ``False``.", stacklevel=2)
    return torch.linalg.cholesky(a, upper=upper, out=out)
//  [python]
def forward(ctx, a):
        l = torch.cholesky(a, False)
        ctx.save_for_backward(l)
        return l
//  [python]
def forward(ctx, a):
        l = torch.linalg.cholesky(a, False)
        ctx.save_for_backward(l)
        return l
//  [python]
def cholesky_inverse(fish, momentum):
    """Performs the inverse of a matrix, using the cholesky inverse (with the vector).

    Parameters
    ----------
    fish : torch.tensor
        Square matrix to be inverted: shape (D,D).
    momentum : torch.tensor
        Vector of shape (D,).

    Returns
    -------
    torch.tensor
        Returns the inverted matrix multiplied by the vector.

    """
    lower = torch.cholesky(fish)
    y = torch.triangular_solve(momentum.view(-1, 1), lower, upper=False, transpose=False, unitriangular=False)[0]
    fish_inv_p = torch.triangular_solve(y, lower.t(), upper=True, transpose=False, unitriangular=False)[0]
    return fish_inv_p
//  [python]
def cholesky_inverse(fish, momentum):
    """Performs the inverse of a matrix, using the cholesky inverse (with the vector).

    Parameters
    ----------
    fish : torch.tensor
        Square matrix to be inverted: shape (D,D).
    momentum : torch.tensor
        Vector of shape (D,).

    Returns
    -------
    torch.tensor
        Returns the inverted matrix multiplied by the vector.

    """
    lower = torch.linalg.cholesky(fish)
    y = torch.triangular_solve(momentum.view(-1, 1), lower, upper=False, transpose=False, unitriangular=False)[0]
    fish_inv_p = torch.triangular_solve(y, lower.t(), upper=True, transpose=False, unitriangular=False)[0]
    return fish_inv_p
//  [python]
def cholesky_solve(mat: np.ndarray, y: np.ndarray) -> np.ndarray:
    if gdml_status["cholesky_engine"] == "torch":
        dev_id = 0

        logger.debug("Using GPU %d to perform PyTorch Cholesky decomposition.", dev_id)
        try:
            cuda_y = torch.tensor(y, device="cuda:%d" % dev_id)
            cuda_mat = torch.tensor(mat, device="cuda:%d" % dev_id)
            u_matrix = torch.cholesky(cuda_mat, upper=True)
            solution = torch.cholesky_solve(cuda_y[:, None], u_matrix, upper=True)
            return solution.cpu().numpy()
        except RuntimeError as e:
            if "CUDA out of memory" in str(e):
                logger.warning(
                    "CUDA is out of memory. Will switch to numpy " "for Cholesky solve."
                )
                gdml_status["cholesky_engine"] = "numpy"
            else:
                raise e

    if gdml_status["cholesky_engine"] == "numpy":
        return sp.linalg.cho_solve(sp.linalg.cho_factor(mat), y)
    else:
        raise ValueError("Unknown linalg backend: %s" % gdml_status["cholesky_engine"])
//  [python]
def cholesky_solve(mat: np.ndarray, y: np.ndarray) -> np.ndarray:
    if gdml_status["cholesky_engine"] == "torch":
        dev_id = 0

        logger.debug("Using GPU %d to perform PyTorch Cholesky decomposition.", dev_id)
        try:
            cuda_y = torch.tensor(y, device="cuda:%d" % dev_id)
            cuda_mat = torch.tensor(mat, device="cuda:%d" % dev_id)
            u_matrix = torch.linalg.cholesky(cuda_mat, upper=True)
            solution = torch.linalg.solve_triangular(u_matrix, cuda_y[:, None], upper=True)
            return solution.cpu().numpy()
        except RuntimeError as e:
            if "CUDA out of memory" in str(e):
                logger.warning(
                    "CUDA is out of memory. Will switch to numpy " "for Cholesky solve."
                )
                gdml_status["cholesky_engine"] = "numpy"
            else:
                raise e

    if gdml_status["cholesky_engine"] == "numpy":
        return sp.linalg.cho_solve(sp.linalg.cho_factor(mat), y)
    else:
        raise ValueError("Unknown linalg backend: %s" % gdml_status["cholesky_engine"])
//  [python]
def update(self, x, P, z, pred_sigma):
        """
        Update state x and P after the observation, 
        outputs filtered state and covariance
        """
        # Update state mean and covariance p(x | y), Joseph Form      
        # Kalman gain, a more stable implementation than naive P @ H^T @ y_sigma^{-1}
        L = torch.cholesky(pred_sigma)
        K = torch.triangular_solve(self.H @ P.transpose(0, 1), L, upper=False)[0]
        K = torch.triangular_solve(K, L.transpose(0, 1))[0].transpose(0, 1)
        
        v = z - self.H @ x
        x = x + K @ v
        P = (torch.eye(*P.shape) - K @ self.H) @ P @ (torch.eye(*P.shape) - K @ self.H).T + K @ self.R @ K.T
        
        assert x.shape == (self.latent_dim, 1)
        assert P.shape == (self.latent_dim, self.latent_dim)
        return x, P
//  [python]
def update(self, x, P, z, pred_sigma):
        """
        Update state x and P after the observation, 
        outputs filtered state and covariance
        """
        # Update state mean and covariance p(x | y), Joseph Form      
        # Kalman gain, a more stable implementation than naive P @ H^T @ y_sigma^{-1}
        L = torch.linalg.cholesky(pred_sigma)
        K = torch.triangular_solve(self.H @ P.transpose(0, 1), L, upper=False)[0]
        K = torch.triangular_solve(K, L.transpose(0, 1))[0].transpose(0, 1)
        
        v = z - self.H @ x
        x = x + K @ v
        P = (torch.eye(*P.shape) - K @ self.H) @ P @ (torch.eye(*P.shape) - K @ self.H).T + K @ self.R @ K.T
        
        assert x.shape == (self.latent_dim, 1)
        assert P.shape == (self.latent_dim, self.latent_dim)
        return x, P
//  [python]
def negELBO(self):
        K_hh = self.kernel.matrix(self.X, torch.exp(self.log_ls_h))
        K_u1 = self.kernel.cross(self.X_col, self.X, torch.exp(self.log_ls_h))
        K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K_u_tt_u_tt = self.kernel.matrix_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_u_tt = self.kernel.cross_pp_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K = torch.cat((torch.cat((K_hh, K_u_tt.T), 1), torch.cat((K_u_tt, K_u_tt_u_tt), 1)), 0)
        A = torch.cholesky(K)
        self.L = torch.diag(self.L1.exp().view(-1)) + torch.matmul(self.L2, self.L2.T)
        self.L = torch.cholesky(self.L)
        Ltril = self.L
        kl = self.KL_term(self.mu.clone().detach(), Ltril, torch.eye(self.M, device=self.device))
        eta = self.mu + torch.matmul(Ltril, torch.randn(self.M, 1, device=self.device))
        f = torch.matmul(A, eta)
        elbo = (-kl + 0.5 * self.N * self.log_tau - 0.5 * torch.exp(self.log_tau) * torch.sum(torch.square(f[0:self.N].view(-1, 1) - self.y)) + 0.5 * self.N_col * self.log_v -
                0.5 * torch.exp(self.log_v) * torch.sum(torch.square(f[self.N_con:self.N_con + self.N_col].view(-1, 1) + torch.sin(f[self.N:self.N_con]).view(-1, 1))))
        return -elbo
//  [python]
def negELBO(self):
        K_hh = self.kernel.matrix(self.X, torch.exp(self.log_ls_h))
        K_u1 = self.kernel.cross(self.X_col, self.X, torch.exp(self.log_ls_h))
        K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K_u_tt_u_tt = self.kernel.matrix_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_u_tt = self.kernel.cross_pp_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K = torch.cat((torch.cat((K_hh, K_u_tt.T), 1), torch.cat((K_u_tt, K_u_tt_u_tt), 1)), 0)
        A = torch.linalg.cholesky(K)
        self.L = torch.diag(self.L1.exp().view(-1)) + torch.matmul(self.L2, self.L2.T)
        self.L = torch.linalg.cholesky(self.L)
        Ltril = self.L
        kl = self.KL_term(self.mu.clone().detach(), Ltril, torch.eye(self.M, device=self.device))
        eta = self.mu + torch.matmul(Ltril, torch.randn(self.M, 1, device=self.device))
        f = torch.matmul(A, eta)
        elbo = (-kl + 0.5 * self.N * self.log_tau - 0.5 * torch.exp(self.log_tau) * torch.sum(torch.square(f[0:self.N].view(-1, 1) - self.y)) + 0.5 * self.N_col * self.log_v -
                0.5 * torch.exp(self.log_v) * torch.sum(torch.square(f[self.N_con:self.N_con + self.N_col].view(-1, 1) + torch.sin(f[self.N:self.N_con]).view(-1, 1))))
        return -elbo
//  [python]
def rsample_logpq_weights(self, XLX, XLY, prior, neuron_prec=True):
    device = XLX.device
    in_features = XLX.shape[-1]
    out_features = XLY.shape[-3]

    assert 4 == len(XLX.shape) and 4 == len(XLY.shape)
    assert XLX.shape[0] == XLY.shape[0]
    assert XLX.shape[-3] in (out_features, 1)
    assert XLY.shape[-1] == 1
    S = XLX.shape[0]
    prior_prec = prior(S)

    prior_prec_full = prior_prec.full()
    if len(prior_prec_full.shape) > 2:
        prior_prec_full = prior_prec_full.unsqueeze(1)

    prec = XLX + prior_prec_full
    L = t.cholesky(prec)

    logdet_prec = 2*L.diagonal(dim1=-1, dim2=-2).log().sum(-1)
    logdet_prec = logdet_prec.expand(S, out_features).sum(-1)

    if self._sample is not None: 
        assert S==self._sample.shape[0]
        dW = self._sample.unsqueeze(-1) - t.cholesky_solve(XLY, L)
        Z = L.transpose(-1, -2) @ dW
    else:
        Z = t.randn(S, out_features, in_features, 1, device=device, dtype=L.dtype)
        dW = t.triangular_solve(Z, L, upper=False, transpose=True)[0]
        self._sample = (t.cholesky_solve(XLY, L) + dW).squeeze(-1)

    logP = mvnormal_log_prob_unnorm(prior_prec, self._sample.transpose(-1, -2))
    logQ = -0.5*(Z**2).sum((-1, -2, -3)) + 0.5*logdet_prec

    logPQw = logP-logQ
    self.logpq = logPQw
    return self._sample
//  [python]
def rsample_logpq_weights(self, XLX, XLY, prior, neuron_prec=True):
    device = XLX.device
    in_features = XLX.shape[-1]
    out_features = XLY.shape[-3]

    assert 4 == len(XLX.shape) and 4 == len(XLY.shape)
    assert XLX.shape[0] == XLY.shape[0]
    assert XLX.shape[-3] in (out_features, 1)
    assert XLY.shape[-1] == 1
    S = XLX.shape[0]
    prior_prec = prior(S)

    prior_prec_full = prior_prec.full()
    if len(prior_prec_full.shape) > 2:
        prior_prec_full = prior_prec_full.unsqueeze(1)

    prec = XLX + prior_prec_full
    L = t.linalg.cholesky(prec)

    logdet_prec = 2*L.diagonal(dim1=-1, dim2=-2).log().sum(-1)
    logdet_prec = logdet_prec.expand(S, out_features).sum(-1)

    if self._sample is not None: 
        assert S==self._sample.shape[0]
        dW = self._sample.unsqueeze(-1) - t.linalg.solve_triangular(XLY, L, upper=False)
        Z = L.transpose(-1, -2) @ dW
    else:
        Z = t.randn(S, out_features, in_features, 1, device=device, dtype=L.dtype)
        dW = t.linalg.solve_triangular(Z, L, upper=False, transpose=True)
        self._sample = (t.linalg.solve_triangular(XLY, L, upper=False) + dW).squeeze(-1)

    logP = mvnormal_log_prob_unnorm(prior_prec, self._sample.transpose(-1, -2))
    logQ = -0.5*(Z**2).sum((-1, -2, -3)) + 0.5*logdet_prec

    logPQw = logP-logQ
    self.logpq = logPQw
    return self._sample
//  [python]
def gauss_newton_loss(self, F0_p, F1_p, F0, pts0, mask):
        noise = (torch.rand_like(pts0)*2-1)
        pts0_noise = pts0 + noise * self.conf.loss.gauss_newton_magnitude

        F0_noise, mask_noise, J = interpolate_tensor(
                F0, pts0_noise, pad=4, return_gradients=True)
        residuals = F0_noise - F1_p

        H = torch.einsum('...dp,...dq->...pq', J, J)
        H = H + torch.eye(H.shape[-1]).to(H) * 1e-6
        g = torch.einsum('...dp,...d->...p', J, residuals)

        H = torch.where(mask[..., None, None], H, torch.eye(H.shape[-1]).to(H))
        H_, g_ = H.cpu(), g.cpu()
        try:
            U = torch.cholesky(H_, upper=True)
        except Exception as e:  # noqa: F841
            import ipdb; ipdb.set_trace()  # noqa: E702

        delta_ = -torch.cholesky_solve(g_[..., None], U, upper=True)[..., 0]
        delta = delta_.to(H)

        pts0_step = pts0_noise + delta
        error = pts0 - pts0_step
        nll = torch.einsum('...p,...pq,...q->...', error, H, error) / 2
        reg = np.log(2*np.pi) - H.logdet() / 2
        nll += self.conf.loss.gauss_newton_reg_weight * reg

        return nll, mask_noise
//  [python]
def gauss_newton_loss(self, F0_p, F1_p, F0, pts0, mask):
        noise = (torch.rand_like(pts0)*2-1)
        pts0_noise = pts0 + noise * self.conf.loss.gauss_newton_magnitude

        F0_noise, mask_noise, J = interpolate_tensor(
                F0, pts0_noise, pad=4, return_gradients=True)
        residuals = F0_noise - F1_p

        H = torch.einsum('...dp,...dq->...pq', J, J)
        H = H + torch.eye(H.shape[-1]).to(H) * 1e-6
        g = torch.einsum('...dp,...d->...p', J, residuals)

        H = torch.where(mask[..., None, None], H, torch.eye(H.shape[-1]).to(H))
        H_, g_ = H.cpu(), g.cpu()
        try:
            U = torch.linalg.cholesky(H_, upper=True)
        except Exception as e:  # noqa: F841
            import ipdb; ipdb.set_trace()  # noqa: E702

        delta_ = -torch.cholesky_solve(g_[..., None], U, upper=True)[..., 0]
        delta = delta_.to(H)

        pts0_step = pts0_noise + delta
        error = pts0 - pts0_step
        nll = torch.einsum('...p,...pq,...q->...', error, H, error) / 2
        reg = np.log(2*np.pi) - H.logdet() / 2
        nll += self.conf.loss.gauss_newton_reg_weight * reg

        return nll, mask_noise
//  [python]
def solve_linear_system(self, A, b, K, delta=0.0, trust_region=False, method='chol'):
    I = torch.eye(self.N, self.N, device=self.device)
    A_t_K = torch.mm(A.t(), K)
    A_t_A = torch.mm(A_t_K, A) 
    if not trust_region:
      LAM = A_t_A + delta*I
    else:
      LAM = A_t_A + delta*(np.diag(np.diag(A_t_A)))
    R = torch.mm(A_t_K, b)
    if method=='chol':
      l = torch.cholesky(LAM, upper=False)
      z = torch.trtrs(R, l, transpose=False, upper=False)[0]
      dtheta = torch.trtrs(z, l, transpose=True, upper=False)[0]
      # dtheta = torch.potrs(R, l, upper=False)
    else:
      raise NotImplementedError
    return dtheta.view(self.num_traj_states, self.state_dim)
//  [python]
def solve_linear_system(self, A, b, K, delta=0.0, trust_region=False, method='chol'):
    I = torch.eye(self.N, self.N, device=self.device)
    A_t_K = torch.mm(A.t(), K)
    A_t_A = torch.mm(A_t_K, A) 
    if not trust_region:
      LAM = A_t_A + delta*I
    else:
      LAM = A_t_A + delta*(np.diag(np.diag(A_t_A)))
    R = torch.mm(A_t_K, b)
    if method=='chol':
      l = torch.linalg.cholesky(LAM, upper=False)
      z = torch.trtrs(R, l, transpose=False, upper=False)[0]
      dtheta = torch.trtrs(z, l, transpose=True, upper=False)[0]
      # dtheta = torch.potrs(R, l, upper=False)
    else:
      raise NotImplementedError
    return dtheta.view(self.num_traj_states, self.state_dim)
//  [python]
def _add_jitter(self, X: Tensor) -> Tensor:
        jitter_prev = 0
        Eye = torch.eye(X.size(-1), device=X.device, dtype=X.dtype).expand(X.shape)
        for i in range(3):
            jitter_new = self._jitter * (10 ** i)
            X = X + (jitter_new - jitter_prev) * Eye
            jitter_prev = jitter_new
            # This may be VERY slow given upstream pytorch issue:
            # https://github.com/pytorch/pytorch/issues/34272
            try:
                _ = torch.cholesky(X)
                warnings.warn(
                    "X is not a p.d. matrix; "
                    f"Added jitter of {jitter_new:.2e} to the diagonal",
                    RuntimeWarning,
                )
                return X
            except RuntimeError:
                continue
        warnings.warn(
            f"Failed to render X p.d. after adding {jitter_new:.2e} jitter",
            RuntimeWarning,
        )
        return X
//  [python]
def _add_jitter(self, X: Tensor) -> Tensor:
        jitter_prev = 0
        Eye = torch.eye(X.size(-1), device=X.device, dtype=X.dtype).expand(X.shape)
        for i in range(3):
            jitter_new = self._jitter * (10 ** i)
            X = X + (jitter_new - jitter_prev) * Eye
            jitter_prev = jitter_new
            # This may be VERY slow given upstream pytorch issue:
            # https://github.com/pytorch/pytorch/issues/34272
            try:
                _ = torch.linalg.cholesky(X)
                warnings.warn(
                    "X is not a p.d. matrix; "
                    f"Added jitter of {jitter_new:.2e} to the diagonal",
                    RuntimeWarning,
                )
                return X
            except RuntimeError:
                continue
        warnings.warn(
            f"Failed to render X p.d. after adding {jitter_new:.2e} jitter",
            RuntimeWarning,
        )
        return X
//  [python]
def pred(self, Xte, cal_llh=False):
        K_u_u = self.kernel.matrix(self.X, torch.exp(self.log_ls_h))
        K_u1 = self.kernel.cross(self.X_col, self.X, torch.exp(self.log_ls_h))
        K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K_utt_utt = self.kernel.matrix_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_utt = self.kernel.cross_pp_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K = torch.cat((torch.cat((K_u_u, K_utt.T), 1), torch.cat((K_utt, K_utt_utt), 1)), 0)
        A = torch.cholesky(K)
        Ltril = torch.tril(self.L)
        eta = self.mu
        f = torch.matmul(A, eta)
        K_u1_te = self.kernel.cross(Xte, self.X, torch.exp(self.log_ls_h))
        K_u2_te = self.kernel.cross(Xte, self.X_col, torch.exp(self.log_ls_h))
        K_utt_te = self.kernel.cross_pp_0(K_u2_te.T, self.X_col, Xte, 0, torch.exp(self.log_ls_h))
        K_te = torch.cat((K_u1_te, K_utt_te.T), 1)
        pred_mean = torch.matmul(K_te, torch.linalg.solve(K, f[0:self.N_con + self.N_col]))
        A1 = torch.cholesky(K)
        t1 = torch.linalg.solve(A1.T, Ltril[0:self.N_con + self.N_col])
        t2 = torch.matmul(t1, t1.T)
        pred_stds = (1 + self.jitter - (K_te * torch.linalg.solve(K, K_te.T).T).sum(1) + (K_te * torch.matmul(t2, K_te.T).T).sum(1))
        pred_stds = pred_stds.view(-1, 1)
        return pred_mean, pred_stds**0.5
//  [python]
def pred(self, Xte, cal_llh=False):
        K_u_u = self.kernel.matrix(self.X, torch.exp(self.log_ls_h))
        K_u1 = self.kernel.cross(self.X_col, self.X, torch.exp(self.log_ls_h))
        K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K_utt_utt = self.kernel.matrix_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_utt = self.kernel.cross_pp_0(K_u1, self.X_col, self.X, 0, torch.exp(self.log_ls_h))
        K = torch.cat((torch.cat((K_u_u, K_utt.T), 1), torch.cat((K_utt, K_utt_utt), 1)), 0)
        A = torch.linalg.cholesky(K)
        Ltril = torch.tril(self.L)
        eta = self.mu
        f = torch.matmul(A, eta)
        K_u1_te = self.kernel.cross(Xte, self.X, torch.exp(self.log_ls_h))
        K_u2_te = self.kernel.cross(Xte, self.X_col, torch.exp(self.log_ls_h))
        K_utt_te = self.kernel.cross_pp_0(K_u2_te.T, self.X_col, Xte, 0, torch.exp(self.log_ls_h))
        K_te = torch.cat((K_u1_te, K_utt_te.T), 1)
        pred_mean = torch.matmul(K_te, torch.linalg.solve(K, f[0:self.N_con + self.N_col]))
        A1 = torch.linalg.cholesky(K)
        t1 = torch.linalg.solve(A1.T, Ltril[0:self.N_con + self.N_col])
        t2 = torch.matmul(t1, t1.T)
        pred_stds = (1 + self.jitter - (K_te * torch.linalg.solve(K, K_te.T).T).sum(1) + (K_te * torch.matmul(t2, K_te.T).T).sum(1))
        pred_stds = pred_stds.view(-1, 1)
        return pred_mean, pred_stds**0.5
//  [python]
def fit(self, x, y):
        self._k_in = self.kernel(x)
        self._L = torch.cholesky(self._k_in)
        self._x = x
        self._y = y
//  [python]
def fit(self, x, y):
        self._k_in = self.kernel(x)
        self._L = torch.linalg.cholesky(self._k_in)
        self._x = x
        self._y = y
//  [python]
def exact_marginal_log_likelihood(kernel_train, training_labels, log_noise_var):
    n = len(training_labels)
    L_train = torch.cholesky(kernel_train + torch.exp(log_noise_var) * torch.eye(n, dtype=torch.float, device=kernel_train.device), upper=False)
    alpha = cholesky_solve(training_labels, L_train)
    mll = -0.5 * training_labels.t().mm(alpha) - L_train.diagonal().log().sum() - (n / 2) * np.log(2*np.pi)

    return mll
//  [python]
def exact_marginal_log_likelihood(kernel_train, training_labels, log_noise_var):
    n = len(training_labels)
    L_train = torch.linalg.cholesky(kernel_train + torch.exp(log_noise_var) * torch.eye(n, dtype=torch.float, device=kernel_train.device))
    alpha = cholesky_solve(training_labels, L_train)
    mll = -0.5 * training_labels.t().mm(alpha) - L_train.diagonal().log().sum() - (n / 2) * np.log(2*np.pi)

    return mll
//  [python]
def cholesky_wrapper(mat, default_diag=None, force_cpu=True):
    device = mat.device
    if force_cpu:
        mat = mat.cpu()
    try:
        tril = torch.cholesky(mat, upper=False)
    except RuntimeError:
        n_dims = mat.size(-1)
        tril = []
        default_tril_single = torch.diag(mat.new_tensor(default_diag)) if default_diag is not None \
            else torch.eye(n_dims, dtype=mat.dtype, device=mat.device)
        for cov in mat.reshape(-1, n_dims, n_dims):
            try:
                tril.append(torch.cholesky(cov, upper=False))
            except RuntimeError:
                tril.append(default_tril_single)
        tril = torch.stack(tril, dim=0).reshape(mat.shape)
    return tril.to(device)
//  [python]
def cholesky_wrapper(mat, default_diag=None, force_cpu=True):
    device = mat.device
    if force_cpu:
        mat = mat.cpu()
    try:
        tril = torch.linalg.cholesky(mat, upper=False)
    except RuntimeError:
        n_dims = mat.size(-1)
        tril = []
        default_tril_single = torch.diag(mat.new_tensor(default_diag)) if default_diag is not None \
            else torch.eye(n_dims, dtype=mat.dtype, device=mat.device)
        for cov in mat.reshape(-1, n_dims, n_dims):
            try:
                tril.append(torch.linalg.cholesky(cov, upper=False))
            except RuntimeError:
                tril.append(default_tril_single)
        tril = torch.stack(tril, dim=0).reshape(mat.shape)
    return tril.to(device)
//  [python]
def test_old_cholesky_empty(self, device, dtype):
        def run_test(upper):
            A = torch.empty(0, 0, dtype=dtype, device=device)
            chol = torch.cholesky(A, upper)
            chol_A = torch.matmul(chol, chol.t().conj())
            self.assertEqual(A, chol_A)
        for upper in [True, False]:
            run_test(upper)
//  [python]
def test_old_cholesky_empty(self, device, dtype):
        def run_test(upper):
            A = torch.empty(0, 0, dtype=dtype, device=device)
            chol = torch.linalg.cholesky(A, upper=upper)
            chol_A = torch.matmul(chol, chol.t().conj())
            self.assertEqual(A, chol_A)
        for upper in [True, False]:
            run_test(upper)
//  [python]
def test_old_cholesky_batched(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def cholesky_test_helper(n, batch_dims, upper):
            A = random_hermitian_pd_matrix(n, *batch_dims, dtype=dtype, device=device)
            cholesky_exp = torch.stack([m.cholesky(upper=upper) for m in A.reshape(-1, n, n)])
            cholesky_exp = cholesky_exp.reshape_as(A)
            self.assertEqual(cholesky_exp, torch.cholesky(A, upper=upper))

        for upper, batchsize in itertools.product([True, False], [(3,), (3, 4), (2, 3, 4)]):
            cholesky_test_helper(3, batchsize, upper)
//  [python]
def test_old_cholesky_batched(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def cholesky_test_helper(n, batch_dims, upper):
            A = random_hermitian_pd_matrix(n, *batch_dims, dtype=dtype, device=device)
            cholesky_exp = torch.stack([m.linalg.cholesky(upper=upper) for m in A.reshape(-1, n, n)])
            cholesky_exp = cholesky_exp.reshape_as(A)
            self.assertEqual(cholesky_exp, torch.linalg.cholesky(A, upper=upper))

        for upper, batchsize in itertools.product([True, False], [(3,), (3, 4), (2, 3, 4)]):
            cholesky_test_helper(3, batchsize, upper)
//  [python]
def test_old_cholesky_batched_many_batches(self, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        def cholesky_test_helper(n, batchsize, device, upper):
            A = random_symmetric_pd_matrix(n, batchsize, dtype=dtype, device=device)
            chol_fact = torch.cholesky(A, upper=upper)
            if upper:
                # Correctness check
                self.assertEqual(A, chol_fact.transpose(-2, -1).matmul(chol_fact))
                # Upper triangular check
                self.assertEqual(chol_fact, chol_fact.triu())
            else:
                # Correctness check
                self.assertEqual(A, chol_fact.matmul(chol_fact.transpose(-2, -1)))
                # Lower triangular check
                self.assertEqual(chol_fact, chol_fact.tril())

        for upper, batchsize in itertools.product([True, False], [262144, 524288]):
            cholesky_test_helper(2, batchsize, device, upper)
//  [python]
def test_old_cholesky_batched_many_batches(self, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        def cholesky_test_helper(n, batchsize, device, upper):
            A = random_symmetric_pd_matrix(n, batchsize, dtype=dtype, device=device)
            chol_fact = torch.linalg.cholesky(A, upper=upper)
            if upper:
                # Correctness check
                self.assertEqual(A, chol_fact.transpose(-2, -1).matmul(chol_fact))
                # Upper triangular check
                self.assertEqual(chol_fact, chol_fact.triu())
            else:
                # Correctness check
                self.assertEqual(A, chol_fact.matmul(chol_fact.transpose(-2, -1)))
                # Lower triangular check
                self.assertEqual(chol_fact, chol_fact.tril())

        for upper, batchsize in itertools.product([True, False], [262144, 524288]):
            cholesky_test_helper(2, batchsize, device, upper)
//  [python]
def test_old_cholesky(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        A = random_hermitian_pd_matrix(10, dtype=dtype, device=device)

        # default Case
        C = torch.cholesky(A)
        B = torch.mm(C, C.t().conj())
        self.assertEqual(A, B, atol=1e-14, rtol=0)

        # test Upper Triangular
        U = torch.cholesky(A, True)
        B = torch.mm(U.t().conj(), U)
        self.assertEqual(A, B, atol=1e-14, rtol=0, msg='cholesky (upper) did not allow rebuilding the original matrix')

        # test Lower Triangular
        L = torch.cholesky(A, False)
        B = torch.mm(L, L.t().conj())
        self.assertEqual(A, B, atol=1e-14, rtol=0, msg='cholesky (lower) did not allow rebuilding the original matrix')