jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, K_train_train, pi):
        A = K_train_train
        D_sqrt = pi.sqrt()
        A.mul_(D_sqrt.unsqueeze(1)).mul_(D_sqrt.unsqueeze(2))
        A.diagonal(dim1=-2, dim2=-1).add_(torch.ones_like(pi))
        torch.cholesky(A, out=A)
        z = A.diagonal(dim1=-2, dim2=-1).log().sum()
        torch.cholesky_solve(DiagLazyTensor(D_sqrt).evaluate(), A, out=A)
        A.mul_(D_sqrt.unsqueeze(2))
        E = A
        return E, z.view(1)
//  [python]
def forward(self, K_train_train, pi):
        A = K_train_train
        D_sqrt = pi.sqrt()
        A.mul_(D_sqrt.unsqueeze(1)).mul_(D_sqrt.unsqueeze(2))
        A.diagonal(dim1=-2, dim2=-1).add_(torch.ones_like(pi))
        A = torch.linalg.cholesky(A)
        z = A.diagonal(dim1=-2, dim2=-1).log().sum()
        torch.linalg.solve_triangular(A, DiagLazyTensor(D_sqrt).evaluate(), upper=False, out=A)
        A.mul_(D_sqrt.unsqueeze(2))
        E = A
        return E, z.view(1)
//  [python]
def forward(self, K):
        for v in [0., 1e-3, 1e-2, 1e-1]:
            try:
                if v > 0:
                    K.diagonal(dim1=-2, dim2=-1).add_(v)
                L = torch.cholesky(K, out=K)
                if v > 0:
                    warnings.warn('K is singular - adding jitter {}'.format(v))
                break
            except RuntimeError as e:
                continue

        return L,
//  [python]
def forward(self, K):
        for v in [0., 1e-3, 1e-2, 1e-1]:
            try:
                if v > 0:
                    K.diagonal(dim1=-2, dim2=-1).add_(v)
                L = torch.linalg.cholesky(K, out=K)
                if v > 0:
                    warnings.warn('K is singular - adding jitter {}'.format(v))
                break
            except RuntimeError as e:
                continue

        return L,
//  [python]
def _compute_aux_llr_1vs1(self):
        I, WV, VV = self._compute_aux_L()
        Lnon = I + VV
        cholLnon = torch.cholesky(Lnon, upper=True)
        mult_icholLnon, logcholLnon = invert_trimat(
            cholLnon, lower=False, right_inv=True, return_logdet=True
        )
        logLnon = 2 * logcholLnon

        Ltar = I + 2 * VV
        cholLtar = torch.cholesky(Ltar, upper=True)
        mult_icholLtar, logcholLtar = invert_trimat(
            cholLtar, lower=False, right_inv=True, return_logdet=True
        )
        logLtar = 2 * logcholLtar

        return WV, mult_icholLnon, mult_icholLtar, logLnon, logLtar
//  [python]
def _compute_aux_llr_1vs1(self):
        I, WV, VV = self._compute_aux_L()
        Lnon = I + VV
        cholLnon = torch.linalg.cholesky(Lnon, upper=True)
        mult_icholLnon, logcholLnon = invert_trimat(
            cholLnon, lower=False, right_inv=True, return_logdet=True
        )
        logLnon = 2 * logcholLnon

        Ltar = I + 2 * VV
        cholLtar = torch.linalg.cholesky(Ltar, upper=True)
        mult_icholLtar, logcholLtar = invert_trimat(
            cholLtar, lower=False, right_inv=True, return_logdet=True
        )
        logLtar = 2 * logcholLtar

        return WV, mult_icholLnon, mult_icholLtar, logLnon, logLtar
//  [python]
def rsample(self, sample_shape=torch.Size()):
        shape = torch.Size(sample_shape) + self.batch_shape
        T = _standard_wishart_tril(self.df, self.event_shape[-1], shape)
        eye = torch.eye(T.shape[-1], dtype=T.dtype, device=T.device).expand_as(T)
        # TODO: Replace with cholesky_solve once its derivative is implemented
        # inv_scale = torch.cholesky_solve(eye, self.scale_tril, upper=False)
        inv_scale_ = torch.triangular_solve(eye, self.scale_tril, upper=False)[0]
        inv_scale = torch.triangular_solve(inv_scale_, self.scale_tril, upper=False, transpose=True)[0]
        assert torch.allclose(inv_scale @ self.scale, eye)
        inv_scale_tril = torch.cholesky(inv_scale, upper=False)
        assert torch.allclose((inv_scale_tril @ inv_scale_tril.transpose(-2, -1)) @ self.scale, eye)
        M = inv_scale_tril @ T
        M_inv = torch.triangular_solve(eye, M, upper=False)[0]
        W = M_inv.transpose(-2, -1) @ M_inv
        assert constraints.lower_cholesky.check(T).all()
        assert constraints.lower_cholesky.check(M).all()
        assert constraints.lower_cholesky.check(M_inv).all()
        assert constraints.positive_definite.check(W).all()
        return W
//  [python]
def rsample(self, sample_shape=torch.Size()):
        shape = torch.Size(sample_shape) + self.batch_shape
        T = _standard_wishart_tril(self.df, self.event_shape[-1], shape)
        eye = torch.eye(T.shape[-1], dtype=T.dtype, device=T.device).expand_as(T)
        # TODO: Replace with cholesky_solve once its derivative is implemented
        # inv_scale = torch.cholesky_solve(eye, self.scale_tril, upper=False)
        inv_scale_ = torch.triangular_solve(eye, self.scale_tril, upper=False)[0]
        inv_scale = torch.triangular_solve(inv_scale_, self.scale_tril, upper=False, transpose=True)[0]
        assert torch.allclose(inv_scale @ self.scale, eye)
        inv_scale_tril = torch.linalg.cholesky(inv_scale, upper=False)
        assert torch.allclose((inv_scale_tril @ inv_scale_tril.transpose(-2, -1)) @ self.scale, eye)
        M = inv_scale_tril @ T
        M_inv = torch.triangular_solve(eye, M, upper=False)[0]
        W = M_inv.transpose(-2, -1) @ M_inv
        assert constraints.lower_cholesky.check(T).all()
        assert constraints.lower_cholesky.check(M).all()
        assert constraints.lower_cholesky.check(M_inv).all()
        assert constraints.positive_definite.check(W).all()
        return W
//  [python]
def forward(self, K_train_train, K_train_test, pi, gp_f_train_map):
        jitters = [0, 1e-3, 1e-2, 1e-1, 1]
        for i, v in enumerate(jitters):
            try:
                K_train_train.diagonal(dim1=-2, dim2=-1).add_(v * torch.ones_like(pi))  # TODO: rewrite this
                L = torch.cholesky(K_train_train, out=K_train_train)
                if i > 0:
                    print('Added jitter of', v)
                break
            except RuntimeError as e:
                if i == len(jitters) - 1:
                    raise e
        alpha = torch.cholesky_solve(gp_f_train_map.squeeze(0).unsqueeze(-1), L.squeeze(0)).squeeze(-1)  # TODO  .squeeze(0) needed?
        if alpha.ndim == 1:
            alpha.unsqueeze_(0)
        mu = contract('cnm,cn->cm', K_train_test, alpha)
        return mu,
//  [python]
def forward(self, K_train_train, K_train_test, pi, gp_f_train_map):
        jitters = [0, 1e-3, 1e-2, 1e-1, 1]
        for i, v in enumerate(jitters):
            try:
                K_train_train.diagonal(dim1=-2, dim2=-1).add_(v * torch.ones_like(pi))  # TODO: rewrite this
                L = torch.linalg.cholesky(K_train_train)
                if i > 0:
                    print('Added jitter of', v)
                break
            except RuntimeError as e:
                if i == len(jitters) - 1:
                    raise e
        alpha = torch.cholesky_solve(gp_f_train_map.squeeze(0).unsqueeze(-1), L.squeeze(0)).squeeze(-1)  # TODO  .squeeze(0) needed?
        if alpha.ndim == 1:
            alpha.unsqueeze_(0)
        mu = contract('cnm,cn->cm', K_train_test, alpha)
        return mu,
//  [python]
def logdet(M):
    # L = torch.potrf(M, upper=False)
    L = torch.cholesky(M)
    return 2 * torch.diag(L).log().sum().unsqueeze(-1).unsqueeze(-1)
//  [python]
def logdet(M):
    L = torch.linalg.cholesky(M)
    return 2 * torch.diag(L).log().sum().unsqueeze(-1).unsqueeze(-1)
//  [python]
def grow_cov(self, beta, lam):
        P = self.P
        Sigma = lam * torch.cholesky_inverse(P)
        Sigma += beta * torch.eye(Sigma.shape[0])
        Pnew = (1./lam) * torch.cholesky_inverse(Sigma)
        Lmat = torch.cholesky(Pnew)
        tril_indices = torch.tril_indices(row=self.d_total, col=self.d_total, offset=0)
        self.L.data = Lmat[tril_indices[0], tril_indices[1]]
//  [python]
def grow_cov(self, beta, lam):
        P = self.P
        Sigma = lam * torch.linalg.cholesky_inverse(P)
        Sigma += beta * torch.eye(Sigma.shape[0])
        Pnew = (1./lam) * torch.linalg.cholesky_inverse(Sigma)
        Lmat = torch.linalg.cholesky(Pnew)
        tril_indices = torch.tril_indices(row=self.d_total, col=self.d_total, offset=0)
        self.L.data = Lmat[tril_indices[0], tril_indices[1]]
//  [python]
def forward(self, Xte):
        n_test = Xte.size(0)
        Xte = (Xte - self.Xmean.expand_as(Xte)) / self.Xstd.expand_as(Xte)

        Xte = self.dkl_nn.get_intermediate_representation(Xte, self.nn_feature_level).view(10, -1)

        X = self.dkl_nn.get_intermediate_representation(self.X, self.nn_feature_level).view(10, -1)

        Sigma = self.kernel(X, X) + self.log_beta.exp().pow(-1) * torch.eye(X.size(0)) \
            + self.jitter * torch.eye(X.size(0))

        kx = self.kernel(X, Xte)
        L = torch.cholesky(Sigma)
        LinvKx, _ = torch.triangular_solve(kx, L, upper=False)

        # option 1
        mean = kx.t() @ torch.cholesky_solve(self.Y, L)  # torch.linalg.cholesky()

        var_diag = self.kernel(Xte, Xte).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim=0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)

        # de-normalized
        mean = mean * self.Ystd.expand_as(mean) + self.Ymean.expand_as(mean)
        var_diag = var_diag.expand_as(mean) * self.Ystd**2

        return mean, var_diag
//  [python]
def forward(self, Xte):
        n_test = Xte.size(0)
        Xte = (Xte - self.Xmean.expand_as(Xte)) / self.Xstd.expand_as(Xte)

        Xte = self.dkl_nn.get_intermediate_representation(Xte, self.nn_feature_level).view(10, -1)

        X = self.dkl_nn.get_intermediate_representation(self.X, self.nn_feature_level).view(10, -1)

        Sigma = self.kernel(X, X) + self.log_beta.exp().pow(-1) * torch.eye(X.size(0)) \
            + self.jitter * torch.eye(X.size(0))

        kx = self.kernel(X, Xte)
        L = torch.linalg.cholesky(Sigma)
        LinvKx, _ = torch.triangular_solve(kx, L, upper=False)

        # option 1
        mean = kx.t() @ torch.linalg.solve_triangular(L, self.Y, upper=False)  # torch.linalg.cholesky()

        var_diag = self.kernel(Xte, Xte).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim=0).view(-1, 1)

        # add the noise uncertainty
        var_diag = var_diag + self.log_beta.exp().pow(-1)

        # de-normalized
        mean = mean * self.Ystd.expand_as(mean) + self.Ymean.expand_as(mean)
        var_diag = var_diag.expand_as(mean) * self.Ystd**2

        return mean, var_diag
//  [python]
def linear_gaussian(
    theta: Tensor, likelihood_shift: Tensor, likelihood_cov: Tensor
) -> Tensor:
    chol_factor = torch.cholesky(likelihood_cov)

    return likelihood_shift + theta + torch.mm(chol_factor, torch.randn_like(theta).T).T
//  [python]
def linear_gaussian(
    theta: Tensor, likelihood_shift: Tensor, likelihood_cov: Tensor
) -> Tensor:
    chol_factor = torch.linalg.cholesky(likelihood_cov)

    return likelihood_shift + theta + torch.mm(chol_factor, torch.randn_like(theta).T).T
//  [python]
def gradientFn(P, lmbda, v):
        """ Compute vector-Jacobian product DJ(M) = DJ(P) DP(M) [b x m*n]

        DP(M) = (H^-1 * A^T * (A * H^-1 * A^T)^-1 * A * H^-1 - H^-1) * B
        H = D_YY^2 f(x, y) = diag(1 / (lmbda * vec(P)))
        B = D_XY^2 f(x, y) = I

        Using:
            Lemma 4.4 from
            Stephen Gould, Richard Hartley, and Dylan Campbell, 2019
            "Deep Declarative Networks: A New Hope", arXiv:1909.04866

        Arguments:
            P: (b, m, n) Torch tensor
                batch of transport matrices

            lmbda: float,
                regularisation factor

            v: (b, m*n) Torch tensor
                batch of gradients of J with respect to P

        Return Values:
            gradient: (b, m*n) Torch tensor,
                batch of gradients of J with respect to M

        """
        with torch.no_grad():
            b, m, n = P.size()
            B = lmbda * P
            hinv = B.flatten(start_dim=-2)
            d1inv = B.sum(-1)[:, 1:].reciprocal() # Remove first element
            d2 = B.sum(-2)
            B = B[:, 1:, :] # Remove top row
            S = -B.transpose(-2, -1).matmul(d1inv.unsqueeze(-1) * B)
            S[:, range(n), range(n)] += d2
            Su = torch.cholesky(S)
            Sinv = torch.zeros_like(S)
            for i in range (b):
                Sinv[i, ...] = torch.cholesky_inverse(Su[i, ...]) # Currently cannot handle batches
            R = -B.matmul(Sinv) * d1inv.unsqueeze(-1)
            Q = -R.matmul(B.transpose(-2, -1)  * d1inv.unsqueeze(-2))
            Q[:, range(m - 1), range(m - 1)] += d1inv
            # Build vector-Jacobian product from left to right:
            vHinv = v * hinv # bxmn * bxmn -> bxmn
            # Break vHinv into m blocks of n elements:
            u1 = vHinv.reshape((-1, m, n)).sum(-1)[:, 1:].unsqueeze(-2) # remove first element
            u2 = vHinv.reshape((-1, m, n)).sum(-2).unsqueeze(-2)
            u3 = u1.matmul(Q) + u2.matmul(R.transpose(-2, -1))
            u4 = u1.matmul(R) + u2.matmul(Sinv)
            u5 = u3.expand(-1, n, -1).transpose(-2, -1) + u4.expand(-1, m - 1, -1)
            uHinv = torch.cat((u4, u5), dim=-2).flatten(start_dim=-2) * hinv
            gradient = uHinv - vHinv
        return gradient
//  [python]
def gradientFn(P, lmbda, v):
        """ Compute vector-Jacobian product DJ(M) = DJ(P) DP(M) [b x m*n]

        DP(M) = (H^-1 * A^T * (A * H^-1 * A^T)^-1 * A * H^-1 - H^-1) * B
        H = D_YY^2 f(x, y) = diag(1 / (lmbda * vec(P)))
        B = D_XY^2 f(x, y) = I

        Using:
            Lemma 4.4 from
            Stephen Gould, Richard Hartley, and Dylan Campbell, 2019
            "Deep Declarative Networks: A New Hope", arXiv:1909.04866

        Arguments:
            P: (b, m, n) Torch tensor
                batch of transport matrices

            lmbda: float,
                regularisation factor

            v: (b, m*n) Torch tensor
                batch of gradients of J with respect to P

        Return Values:
            gradient: (b, m*n) Torch tensor,
                batch of gradients of J with respect to M

        """
        with torch.no_grad():
            b, m, n = P.size()
            B = lmbda * P
            hinv = B.flatten(start_dim=-2)
            d1inv = B.sum(-1)[:, 1:].reciprocal() # Remove first element
            d2 = B.sum(-2)
            B = B[:, 1:, :] # Remove top row
            S = -B.transpose(-2, -1).matmul(d1inv.unsqueeze(-1) * B)
            S[:, range(n), range(n)] += d2
            Su = torch.linalg.cholesky(S)
            Sinv = torch.zeros_like(S)
            for i in range (b):
                Sinv[i, ...] = torch.linalg.cholesky_inverse(Su[i, ...]) # Currently cannot handle batches
            R = -B.matmul(Sinv) * d1inv.unsqueeze(-1)
            Q = -R.matmul(B.transpose(-2, -1)  * d1inv.unsqueeze(-2))
            Q[:, range(m - 1), range(m - 1)] += d1inv
            # Build vector-Jacobian product from left to right:
            vHinv = v * hinv # bxmn * bxmn -> bxmn
            # Break vHinv into m blocks of n elements:
            u1 = vHinv.reshape((-1, m, n)).sum(-1)[:, 1:].unsqueeze(-2) # remove first element
            u2 = vHinv.reshape((-1, m, n)).sum(-2).unsqueeze(-2)
            u3 = u1.matmul(Q) + u2.matmul(R.transpose(-2, -1))
            u4 = u1.matmul(R) + u2.matmul(Sinv)
            u5 = u3.expand(-1, n, -1).transpose(-2, -1) + u4.expand(-1, m - 1, -1)
            uHinv = torch.cat((u4, u5), dim=-2).flatten(start_dim=-2) * hinv
            gradient = uHinv - vHinv
        return gradient
//  [python]
def forward(
        self, *, observations: types.ObservationsTorch
    ) -> Tuple[types.StatesTorch, types.ScaleTrilTorch]:
        """Predicts states and uncertainties from observation inputs.

        Uncertainties should be lower-triangular Cholesky decompositions of covariance
        matrices.

        Args:
            observations (dict or torch.Tensor): Measurement inputs. Should be
                either a dict of tensors or tensor of size `(N, ...)`.
        Returns:
            Tuple[torch.Tensor, torch.Tensor]: Predicted states & uncertainties.
                - States should have shape `(N, state_dim).`
                - Uncertainties should be lower triangular, and should have shape
                `(N, state_dim, state_dim).`
        """
        # Observations should be tensor, not dictionary
        assert isinstance(observations, torch.Tensor)
        observations = cast(torch.Tensor, observations)
        N = observations.shape[0]

        # Compute/return predicted state and uncertainty values
        # Note that for square C_pinv matrices, we can compute scale_tril as C_pinv @
        # R_tril. In the general case, we transform the full covariance and then take
        # the cholesky decomposition.
        predicted_states = (C_pinv[None, :, :] @ observations[:, :, None]).squeeze(-1)
        scale_tril = torch.cholesky(
            C_pinv @ R_tril @ R_tril.transpose(-1, -2) @ C_pinv.transpose(-1, -2)
        )[None, :, :].expand((N, state_dim, state_dim))

        # Add output bias if trainable
        if self.trainable:
            predicted_states += self.output_bias

        return predicted_states, scale_tril
//  [python]
def forward(
        self, *, observations: types.ObservationsTorch
    ) -> Tuple[types.StatesTorch, types.ScaleTrilTorch]:
        """Predicts states and uncertainties from observation inputs.

        Uncertainties should be lower-triangular Cholesky decompositions of covariance
        matrices.

        Args:
            observations (dict or torch.Tensor): Measurement inputs. Should be
                either a dict of tensors or tensor of size `(N, ...)`.
        Returns:
            Tuple[torch.Tensor, torch.Tensor]: Predicted states & uncertainties.
                - States should have shape `(N, state_dim).`
                - Uncertainties should be lower triangular, and should have shape
                `(N, state_dim, state_dim).`
        """
        # Observations should be tensor, not dictionary
        assert isinstance(observations, torch.Tensor)
        observations = cast(torch.Tensor, observations)
        N = observations.shape[0]

        # Compute/return predicted state and uncertainty values
        # Note that for square C_pinv matrices, we can compute scale_tril as C_pinv @
        # R_tril. In the general case, we transform the full covariance and then take
        # the cholesky decomposition.
        predicted_states = (C_pinv[None, :, :] @ observations[:, :, None]).squeeze(-1)
        scale_tril = torch.linalg.cholesky(
            C_pinv @ R_tril @ R_tril.transpose(-1, -2) @ C_pinv.transpose(-1, -2)
        )[None, :, :].expand((N, state_dim, state_dim))

        # Add output bias if trainable
        if self.trainable:
            predicted_states += self.output_bias

        return predicted_states, scale_tril
//  [python]
def test_update_cholesky():
    """Test that the update cholesky function returns correct values."""
    n = 6
    new_A = torch.rand(n, n, dtype=torch.float64)
    new_A = new_A @ new_A.t()
    new_A += torch.eye(len(new_A), dtype=torch.float64)

    A = new_A[:n - 1, :n - 1]

    old_chol = torch.cholesky(A, upper=False)
    new_row = new_A[-1]

    # Test updateing overall
    new_chol = update_cholesky(old_chol, new_row)
    error = new_chol - torch.cholesky(new_A, upper=False)
    assert torch.all(torch.abs(error) <= 1e-15)

    # Test updating inplace
    new_chol = torch.zeros(n, n, dtype=torch.float64)
    new_chol[:n - 1, :n - 1] = old_chol

    update_cholesky(old_chol, new_row, chol_row_out=new_chol[-1])
    error = new_chol - torch.cholesky(new_A, upper=False)
    assert torch.all(torch.abs(error) <= 1e-15)
//  [python]
def test_update_cholesky():
    """Test that the update cholesky function returns correct values."""
    n = 6
    new_A = torch.rand(n, n, dtype=torch.float64)
    new_A = new_A @ new_A.t()
    new_A += torch.eye(len(new_A), dtype=torch.float64)

    A = new_A[:n - 1, :n - 1]

    old_chol = torch.linalg.cholesky(A, upper=False)
    new_row = new_A[-1]

    # Test updateing overall
    new_chol = update_cholesky(old_chol, new_row)
    error = new_chol - torch.linalg.cholesky(new_A, upper=False)
    assert torch.all(torch.abs(error) <= 1e-15)

    # Test updating inplace
    new_chol = torch.zeros(n, n, dtype=torch.float64)
    new_chol[:n - 1, :n - 1] = old_chol

    update_cholesky(old_chol, new_row, chol_row_out=new_chol[-1])
    error = new_chol - torch.linalg.cholesky(new_A, upper=False)
    assert torch.all(torch.abs(error) <= 1e-15)
//  [python]
def cholesky(device_config, K, Y):
    """
    Solve linear system using a cholesky solver.
    Params:
        K - Covariance Matrix
        Y - Target Labels
    """

    with torch.no_grad():
        L = torch.cholesky(K, upper=False)
        solution = torch.cholesky_solve(Y, L, upper=False)

    return solution
//  [python]
def cholesky(device_config, K, Y):
    """
    Solve linear system using a cholesky solver.
    Params:
        K - Covariance Matrix
        Y - Target Labels
    """

    with torch.no_grad():
        L = torch.linalg.cholesky(K, upper=False)
        solution = torch.linalg.solve_triangular(Y, L, upper=False)

    return solution
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None, max_tries=3):
    """Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
    Args:
        :attr:`A` (Tensor):
            The tensor to compute the Cholesky decomposition of
        :attr:`upper` (bool, optional):
            See torch.cholesky
        :attr:`out` (Tensor, optional):
            See torch.cholesky
        :attr:`jitter` (float, optional):
            The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen
            as 1e-6 (float) or 1e-8 (double)
        :attr:`max_tries` (int, optional):
            Number of attempts (with successively increasing jitter) to make before raising an error.
    """
    try:
        L = torch.cholesky(A, upper=upper, out=out)
        return L
    except RuntimeError as e:
        isnan = torch.isnan(A)
        if isnan.any():
            raise RuntimeError(
                f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
            )

        if jitter is None:
            jitter = 1e-6 if A.dtype == torch.float32 else 1e-8
        Aprime = A.clone()
        jitter_prev = 0
        for i in range(max_tries):
            jitter_new = jitter * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                L = torch.cholesky(Aprime, upper=upper, out=out)
                warnings.warn(f"A not p.d., added jitter of {jitter_new:.1e} to the diagonal", RuntimeWarning)
                return L
            except RuntimeError:
                continue
        raise CholeskyError(
            f"Matrix not positive definite after repeatedly adding jitter up to {jitter_new:.1e}. "
            f"Original error on first attempt: {e}"
        )
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None, max_tries=3):
    """Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
    Args:
        :attr:`A` (Tensor):
            The tensor to compute the Cholesky decomposition of
        :attr:`upper` (bool, optional):
            See torch.linalg.cholesky
        :attr:`out` (Tensor, optional):
            See torch.linalg.cholesky
        :attr:`jitter` (float, optional):
            The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen
            as 1e-6 (float) or 1e-8 (double)
        :attr:`max_tries` (int, optional):
            Number of attempts (with successively increasing jitter) to make before raising an error.
    """
    try:
        L = torch.linalg.cholesky(A, upper=upper, out=out)
        return L
    except RuntimeError as e:
        isnan = torch.isnan(A)
        if isnan.any():
            raise RuntimeError(
                f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
            )

        if jitter is None:
            jitter = 1e-6 if A.dtype == torch.float32 else 1e-8
        Aprime = A.clone()
        jitter_prev = 0
        for i in range(max_tries):
            jitter_new = jitter * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                L = torch.linalg.cholesky(Aprime, upper=upper, out=out)
                warnings.warn(f"A not p.d., added jitter of {jitter_new:.1e} to the diagonal", RuntimeWarning)
                return L
            except RuntimeError:
                continue
        raise CholeskyError(
            f"Matrix not positive definite after repeatedly adding jitter up to {jitter_new:.1e}. "
            f"Original error on first attempt: {e}"
        )
//  [python]
def cholesky_solve(ATA, ATy):
    if ATA.dtype == torch.half or ATy.dtype == torch.half:
        return ATy.to(torch.float).cholesky_solve(torch.cholesky(ATA.to(torch.float))).to(ATy.dtype)
    return ATy.cholesky_solve(torch.cholesky(ATA)).to(ATy.dtype)
//  [python]
def cholesky_solve(ATA, ATy):
    if ATA.dtype == torch.half or ATy.dtype == torch.half:
        return ATy.to(torch.float).cholesky_solve(torch.linalg.cholesky(ATA.to(torch.float))).to(ATy.dtype)
    return ATy.cholesky_solve(torch.linalg.cholesky(ATA)).to(ATy.dtype)
//  [python]
def cholesky(M, eps=1e-4):
  '''
  Compute L, where M = LL^T.
  '''
  I = torch.eye(M.size(-1), device=M.device)
  L = torch.cholesky(M + eps * I , upper=False)
  return L
//  [python]
def cholesky(M, eps=1e-4):
  '''
  Compute L, where M = LL^T.
  '''
  I = torch.eye(M.size(-1), device=M.device)
  L = torch.linalg.cholesky(M + eps * I, upper=False)
  return L
//  [python]
def scatter_encoding(self, F, y, mask, sigma_2=1e-6):
        # B, M, N
        # B, c, N
        # B, N
        B,M,N = F.shape
        y_T = y.transpose(1,2)
        F_T = F.transpose(1,2)
        st = time()
        I_N = torch.eye(N).unsqueeze(0).to(F)
        K = (F_T*self.eigenvalues_11k).bmm(F)  # B,N,N
        # set unmasked diagonal as 1 
        masked_diag = I_N.repeat(B,1,1) # B,N,N
        masked_diag[mask,:] *= sigma_2
        K += masked_diag
        u = torch.cholesky(K)
        del F_T,y,I_N,K,masked_diag,mask
        K_inv_y = torch.cholesky_solve(y_T, u)
        #K_inv_y = torch.linalg.inv(K).bmm(y_T)
        latent = (F*self.eigenvalues_1k1).bmm(K_inv_y)
        del F, K_inv_y, y_T, u
        torch.cuda.empty_cache()
        return latent
//  [python]
def scatter_encoding(self, F, y, mask, sigma_2=1e-6):
        # B, M, N
        # B, c, N
        # B, N
        B,M,N = F.shape
        y_T = y.transpose(1,2)
        F_T = F.transpose(1,2)
        st = time()
        I_N = torch.eye(N).unsqueeze(0).to(F)
        K = (F_T*self.eigenvalues_11k).bmm(F)  # B,N,N
        # set unmasked diagonal as 1 
        masked_diag = I_N.repeat(B,1,1) # B,N,N
        masked_diag[mask,:] *= sigma_2
        K += masked_diag
        u = torch.linalg.cholesky(K)
        del F_T,y,I_N,K,masked_diag,mask
        K_inv_y = torch.cholesky_solve(y_T, u)
        #K_inv_y = torch.linalg.inv(K).bmm(y_T)
        latent = (F*self.eigenvalues_1k1).bmm(K_inv_y)
        del F, K_inv_y, y_T, u
        torch.cuda.empty_cache()
        return latent
//  [python]
def forward(self, x):
        """
        The RHVAE model
        """

        recon_x, z0, eps0, mu, log_var = self.vae_forward(x)

        z = z0

        if self.training:

            # update the metric using batch data points
            L, M = self.metric_forward(x)

            # store LL^T and mu(x_i) to update final metric
            self.M.append(M.clone().detach())
            self.centroids.append(mu.clone().detach())

            G_inv = (
                M.unsqueeze(0)
                * torch.exp(
                    -torch.norm(mu.unsqueeze(0) - z.unsqueeze(1), dim=-1) ** 2
                    / (self.T ** 2)
                )
                .unsqueeze(-1)
                .unsqueeze(-1)
            ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(self.device)

        else:
            G = self.G(z)
            G_inv = self.G_inv(z)
            L = torch.cholesky(G)

        G_log_det = -torch.logdet(G_inv)

        gamma = torch.randn_like(z0, device=self.device)
        rho = gamma / self.beta_zero_sqrt
        beta_sqrt_old = self.beta_zero_sqrt

        # sample \rho from N(0, G)
        rho = (L @ rho.unsqueeze(-1)).squeeze(-1)

        recon_x = self.decode(z)

        for k in range(self.n_lf):

            # perform leapfrog steps

            # step 1
            rho_ = self.leap_step_1(recon_x, x, z, rho, G_inv, G_log_det)

            # step 2
            z = self.leap_step_2(recon_x, x, z, rho_, G_inv, G_log_det)

            recon_x = self.decode(z)

            if self.training:
                G_inv = (
                    M.unsqueeze(0)
                    * torch.exp(
                        -torch.norm(mu.unsqueeze(0) - z.unsqueeze(1), dim=-1) ** 2
                        / (self.T ** 2)
                    )
                    .unsqueeze(-1)
                    .unsqueeze(-1)
                ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(self.device)

            else:
                # compute metric value on new z using final metric
                G = self.G(z)
                G_inv = self.G_inv(z)

            G_log_det = -torch.logdet(G_inv)

            # step 3
            rho__ = self.leap_step_3(recon_x, x, z, rho_, G_inv, G_log_det)

            # tempering
            beta_sqrt = self._tempering(k + 1, self.n_lf)
            rho = (beta_sqrt_old / beta_sqrt) * rho__
            beta_sqrt_old = beta_sqrt

        return recon_x, z, z0, rho, eps0, gamma, mu, log_var, G_inv, G_log_det
//  [python]
def forward(self, x):
        """
        The RHVAE model
        """

        recon_x, z0, eps0, mu, log_var = self.vae_forward(x)

        z = z0

        if self.training:

            # update the metric using batch data points
            L, M = self.metric_forward(x)

            # store LL^T and mu(x_i) to update final metric
            self.M.append(M.clone().detach())
            self.centroids.append(mu.clone().detach())

            G_inv = (
                M.unsqueeze(0)
                * torch.exp(
                    -torch.norm(mu.unsqueeze(0) - z.unsqueeze(1), dim=-1) ** 2
                    / (self.T ** 2)
                )
                .unsqueeze(-1)
                .unsqueeze(-1)
            ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(self.device)

        else:
            G = self.G(z)
            G_inv = self.G_inv(z)
            L = torch.linalg.cholesky(G)

        G_log_det = -torch.logdet(G_inv)

        gamma = torch.randn_like(z0, device=self.device)
        rho = gamma / self.beta_zero_sqrt
        beta_sqrt_old = self.beta_zero_sqrt

        # sample \rho from N(0, G)
        rho = (L @ rho.unsqueeze(-1)).squeeze(-1)

        recon_x = self.decode(z)

        for k in range(self.n_lf):

            # perform leapfrog steps

            # step 1
            rho_ = self.leap_step_1(recon_x, x, z, rho, G_inv, G_log_det)

            # step 2
            z = self.leap_step_2(recon_x, x, z, rho_, G_inv, G_log_det)

            recon_x = self.decode(z)

            if self.training:
                G_inv = (
                    M.unsqueeze(0)
                    * torch.exp(
                        -torch.norm(mu.unsqueeze(0) - z.unsqueeze(1), dim=-1) ** 2
                        / (self.T ** 2)
                    )
                    .unsqueeze(-1)
                    .unsqueeze(-1)
                ).sum(dim=1) + self.lbd * torch.eye(self.latent_dim).to(self.device)

            else:
                # compute metric value on new z using final metric
                G = self.G(z)
                G_inv = self.G_inv(z)

            G_log_det = -torch.logdet(G_inv)

            # step 3
            rho__ = self.leap_step_3(recon_x, x, z, rho_, G_inv, G_log_det)

            # tempering
            beta_sqrt = self._tempering(k + 1, self.n_lf)
            rho = (beta_sqrt_old / beta_sqrt) * rho__
            beta_sqrt_old = beta_sqrt

        return recon_x, z, z0, rho, eps0, gamma, mu, log_var, G_inv, G_log_det
//  [python]
def posdef_logdet(A: torch.Tensor) -> (torch.Tensor, torch.Tensor):
    tril = torch.cholesky(A, upper=False)
    return 2. * triangular_logdet(tril), tril
//  [python]
def posdef_logdet(A: torch.Tensor) -> (torch.Tensor, torch.Tensor):
    tril = torch.linalg.cholesky(A, upper=False)
    return 2. * triangular_logdet(tril), tril
//  [python]
def _sample_approx_posterior_mog(
        self, num_samples, x: Tensor, batch_size: int
    ) -> Tensor:
        r"""
        Sample from the approximate posterior.

        Args:
            num_samples: Desired number of samples.
            x: Conditioning context for posterior $p(\theta|x)$.
            batch_size: Batch size for sampling.

        Returns:
            Samples from the approximate mixture of Gaussians posterior.
        """

        # Compute the mixture components of the posterior.
        logits_p, m_p, prec_p = self._posthoc_correction(x)

        # Compute the precision factors which represent the upper triangular matrix
        # of the cholesky decomposition of the prec_p.
        prec_factors_p = torch.cholesky(prec_p, upper=True)

        assert logits_p.ndim == 2
        assert m_p.ndim == 3
        assert prec_p.ndim == 4
        assert prec_factors_p.ndim == 4

        # Replicate to use batched sampling from pyknos.
        if batch_size is not None and batch_size > 1:
            logits_p = logits_p.repeat(batch_size, 1)
            m_p = m_p.repeat(batch_size, 1, 1)
            prec_factors_p = prec_factors_p.repeat(batch_size, 1, 1, 1)

        # Get (optionally z-scored) MoG samples.
        theta = MultivariateGaussianMDN.sample_mog(
            num_samples, logits_p, m_p, prec_factors_p
        )

        embedded_context = self._neural_net._embedding_net(x)
        if embedded_context is not None:
            # Merge the context dimension with sample dimension in order to
            # apply the transform.
            theta = torchutils.merge_leading_dims(theta, num_dims=2)
            embedded_context = torchutils.repeat_rows(
                embedded_context, num_reps=num_samples
            )

        theta, _ = self._neural_net._transform.inverse(theta, context=embedded_context)

        if embedded_context is not None:
            # Split the context dimension from sample dimension.
            theta = torchutils.split_leading_dim(theta, shape=[-1, num_samples])

        return theta
//  [python]
def _sample_approx_posterior_mog(
        self, num_samples, x: Tensor, batch_size: int
    ) -> Tensor:
        r"""
        Sample from the approximate posterior.

        Args:
            num_samples: Desired number of samples.
            x: Conditioning context for posterior $p(\theta|x)$.
            batch_size: Batch size for sampling.

        Returns:
            Samples from the approximate mixture of Gaussians posterior.
        """

        # Compute the mixture components of the posterior.
        logits_p, m_p, prec_p = self._posthoc_correction(x)

        # Compute the precision factors which represent the upper triangular matrix
        # of the cholesky decomposition of the prec_p.
        prec_factors_p = torch.linalg.cholesky(prec_p, upper=True)

        assert logits_p.ndim == 2
        assert m_p.ndim == 3
        assert prec_p.ndim == 4
        assert prec_factors_p.ndim == 4

        # Replicate to use batched sampling from pyknos.
        if batch_size is not None and batch_size > 1:
            logits_p = logits_p.repeat(batch_size, 1)
            m_p = m_p.repeat(batch_size, 1, 1)
            prec_factors_p = prec_factors_p.repeat(batch_size, 1, 1, 1)

        # Get (optionally z-scored) MoG samples.
        theta = MultivariateGaussianMDN.sample_mog(
            num_samples, logits_p, m_p, prec_factors_p
        )

        embedded_context = self._neural_net._embedding_net(x)
        if embedded_context is not None:
            # Merge the context dimension with sample dimension in order to
            # apply the transform.
            theta = torchutils.merge_leading_dims(theta, num_dims=2)
            embedded_context = torchutils.repeat_rows(
                embedded_context, num_reps=num_samples
            )

        theta, _ = self._neural_net._transform.inverse(theta, context=embedded_context)

        if embedded_context is not None:
            # Split the context dimension from sample dimension.
            theta = torchutils.split_leading_dim(theta, shape=[-1, num_samples])

        return theta
//  [python]
def inverse_cholesky(A: torch.Tensor) -> torch.Tensor:
    r"""Computes the Cholesky decomposition of the inverse: :math:`A^{-1} = LL^\top`.

    Adapted from torch.distributions.multivariate_normal, inspired by:

    - https://nbviewer.jupyter.org/gist/fehiepsi/5ef8e09e61604f10607380467eb82006
    - https://math.stackexchange.com/q/1434899
    """
    Lf = torch.cholesky(A.flip(-2, -1), upper=False)
    L_inv = Lf.flip(-2, -1).transpose(-2, -1)
    L = torch.triangular_solve(eye_like(A), L_inv, upper=False)[0]
    return L
//  [python]
def inverse_cholesky(A: torch.Tensor) -> torch.Tensor:
    r"""Computes the Cholesky decomposition of the inverse: :math:`A^{-1} = LL^\top`.

    Adapted from torch.distributions.multivariate_normal, inspired by:

    - https://nbviewer.jupyter.org/gist/fehiepsi/5ef8e09e61604f10607380467eb82006
    - https://math.stackexchange.com/q/1434899
    """
    Lf = torch.linalg.cholesky(A.flip(-2, -1))
    L_inv = Lf.flip(-2, -1).transpose(-2, -1)
    L = torch.triangular_solve(eye_like(A), L_inv, upper=False)[0]
    return L
//  [python]
def posdef_solve(b: torch.Tensor, A: torch.Tensor) -> (torch.Tensor, torch.Tensor):
    r"""Computes batched :math:`A^{-1} b`, assuming :math:`A` is positive definite.

    Adapted from: https://github.com/pytorch/pytorch/issues/4669#issuecomment-520919549
    """
    tril = torch.cholesky(A, upper=False)
    x = torch.triangular_solve(b, tril, upper=False)[0]
    return torch.triangular_solve(x, tril, upper=False, transpose=True)[0], tril
//  [python]
def posdef_solve(b: torch.Tensor, A: torch.Tensor) -> (torch.Tensor, torch.Tensor):
    r"""Computes batched :math:`A^{-1} b`, assuming :math:`A` is positive definite.

    Adapted from: https://github.com/pytorch/pytorch/issues/4669#issuecomment-520919549
    """
    tril = torch.linalg.cholesky(A, upper=False)
    x = torch.triangular_solve(b, tril, upper=False)[0]
    return torch.triangular_solve(x, tril, upper=False, transpose=True)[0], tril
//  [python]
def mahalanobis(A, b):
    r"""Computes batched :math:`b^\top A^{-1} b`, assuming :math:`A` is positive definite."""
    tril = torch.cholesky(A, upper=False)
    x = torch.triangular_solve(b.unsqueeze(-1), tril, upper=False)[0].squeeze(-1)
    return (x * x).sum(-1), tril
//  [python]
def mahalanobis(A, b):
    r"""Computes batched :math:`b^\top A^{-1} b`, assuming :math:`A` is positive definite."""
    tril = torch.linalg.cholesky(A, upper=False)
    x = torch.triangular_solve(b.unsqueeze(-1), tril, upper=False)[0].squeeze(-1)
    return (x * x).sum(-1), tril
//  [python]
def compute_damped_gn_update(jacobian, output_error, damping):
    """
    Compute the damped Gauss-Newton update, based on the given jacobian and
    output error.
    Args:
        jacobian (torch.Tensor): 2D tensor containing the Jacobian of the
            flattened output with respect to the flattened parameters for which
            the GN update is computed.
        output_error (torch.Tensor): tensor containing the gradient of the loss
            with respect to the output layer of the network.
        damping (float): positive damping hyperparameter
    Returns: the damped Gauss-Newton update for the parameters for which the
        jacobian was computed.
    """

    if damping < 0:
        raise ValueError('Positive value for damping expected, got '
                         '{}'.format(damping))
    output_error = output_error.view(-1, 1).detach()

    if damping == 0:
        jacobian_pinv = torch.pinverse(jacobian)
        gn_updates = jacobian_pinv.mm(output_error)
    else:
        if jacobian.shape[0] >= jacobian.shape[1]:
            G = jacobian.t().mm(jacobian)
            C = G + damping * torch.eye(G.shape[0])
            C_cholesky = torch.cholesky(C)
            jacobian_error = jacobian.t().matmul(output_error)
            gn_updates = torch.cholesky_solve(jacobian_error, C_cholesky)
        else:
            G = jacobian.mm(jacobian.t())
            C = G + damping * torch.eye(G.shape[0])
            C_cholesky = torch.cholesky(C)
            inverse_error = torch.cholesky_solve(output_error, C_cholesky)
            gn_updates = jacobian.t().matmul(inverse_error)

    return gn_updates
//  [python]
def compute_damped_gn_update(jacobian, output_error, damping):
    """
    Compute the damped Gauss-Newton update, based on the given jacobian and
    output error.
    Args:
        jacobian (torch.Tensor): 2D tensor containing the Jacobian of the
            flattened output with respect to the flattened parameters for which
            the GN update is computed.
        output_error (torch.Tensor): tensor containing the gradient of the loss
            with respect to the output layer of the network.
        damping (float): positive damping hyperparameter
    Returns: the damped Gauss-Newton update for the parameters for which the
        jacobian was computed.
    """

    if damping < 0:
        raise ValueError('Positive value for damping expected, got '
                         '{}'.format(damping))
    output_error = output_error.view(-1, 1).detach()

    if damping == 0:
        jacobian_pinv = torch.pinverse(jacobian)
        gn_updates = jacobian_pinv.mm(output_error)
    else:
        if jacobian.shape[0] >= jacobian.shape[1]:
            G = jacobian.t().mm(jacobian)
            C = G + damping * torch.eye(G.shape[0])
            C_cholesky = torch.linalg.cholesky(C)
            jacobian_error = jacobian.t().matmul(output_error)
            gn_updates = torch.cholesky_solve(jacobian_error, C_cholesky)
        else:
            G = jacobian.mm(jacobian.t())
            C = G + damping * torch.eye(G.shape[0])
            C_cholesky = torch.linalg.cholesky(C)
            inverse_error = torch.cholesky_solve(output_error, C_cholesky)
            gn_updates = jacobian.t().matmul(inverse_error)

    return gn_updates
//  [python]
def __init__(self, mu, precision):
        # mu: [batch_size, z_dim]
        self.mu = mu
        # precision: [batch_size, z_dim, z_dim]
        self.precision = precision
        # TODO: get rid of the inverse for efficiency
        self.L = torch.cholesky(torch.inverse(precision))
        self.dim = self.mu.shape[1]
//  [python]
def __init__(self, mu, precision):
        # mu: [batch_size, z_dim]
        self.mu = mu
        # precision: [batch_size, z_dim, z_dim]
        self.precision = precision
        # TODO: get rid of the inverse for efficiency
        self.L = torch.linalg.cholesky(torch.inverse(precision))
        self.dim = self.mu.shape[1]
//  [python]
def test_to_pos(self):
        for _ in range(1000):
            A_diag = LA.to_positive(torch.randn(10))
            self.assertTrue((A_diag > 0).all())

            try:
                torch.cholesky(torch.diag(A_diag))
            except RuntimeError:
                self.fail('Call to `torch.cholesky` failed.')
//  [python]
def test_to_pos(self):
        for _ in range(1000):
            A_diag = LA.to_positive(torch.randn(10))
            self.assertTrue((A_diag > 0).all())

            try:
                torch.linalg.cholesky(torch.diag(A_diag))
            except RuntimeError:
                self.fail('Call to `torch.cholesky` failed.')
//  [python]
def cholesky_update(self, hyper):
		self.gram_mat_update(hyper)
		eye_mat = torch.diag(self.gram_mat.new_ones(self.gram_mat.size(0)))
		for jitter_const in [0, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3]:
			chol_jitter = torch.trace(self.gram_mat).item() * jitter_const
			try:
				# cholesky is lower triangular matrix
				self.cholesky = torch.cholesky(self.gram_mat + eye_mat * chol_jitter, upper=False)
				self.jitter = chol_jitter
				return
			except RuntimeError:
				pass
		raise RuntimeError('Absolute entry values of Gram matrix are between %.4E~%.4E with trace %.4E' %
		                   (torch.min(torch.abs(self.gram_mat)).item(), torch.max(torch.abs(self.gram_mat)).item(),
		                    torch.trace(self.gram_mat).item()))
//  [python]
def cholesky_update(self, hyper):
		self.gram_mat_update(hyper)
		eye_mat = torch.diag(self.gram_mat.new_ones(self.gram_mat.size(0)))
		for jitter_const in [0, 1e-8, 1e-7, 1e-6, 1e-5, 1e-4, 1e-3]:
			chol_jitter = torch.trace(self.gram_mat).item() * jitter_const
			try:
				# cholesky is lower triangular matrix
				self.cholesky = torch.linalg.cholesky(self.gram_mat + eye_mat * chol_jitter, upper=False)
				self.jitter = chol_jitter
				return
			except RuntimeError:
				pass
		raise RuntimeError('Absolute entry values of Gram matrix are between %.4E~%.4E with trace %.4E' %
		                   (torch.min(torch.abs(self.gram_mat)).item(), torch.max(torch.abs(self.gram_mat)).item(),
		                    torch.trace(self.gram_mat).item()))
//  [python]
def batchedDet_old(batchedTensor):
    jitter = 1e-6
    if np.prod(batchedTensor.shape[:-2]) >= 256 * 256 - 1:
        chunk_num = int(np.prod(batchedTensor.shape[1:-2]))
        if chunk_num >= (256*256 - 1):
            print("TOO BIG TENSOR")
        max_split = (256 * 256 - 1)//chunk_num
        det_list = []
        for t in torch.split(batchedTensor, max_split):
            temp.append(torch.prod(torch.diagonal(torch.cholesky(t), dim1=-2, dim2=-1), dim = -1)**2)
        return torch.cat(temp)
    else:
        return torch.prod(torch.diagonal(torch.cholesky(batchedTensor), dim1=-2, dim2=-1), dim = -1)**2
//  [python]
def batchedDet_old(batchedTensor):
    jitter = 1e-6
    if np.prod(batchedTensor.shape[:-2]) >= 256 * 256 - 1:
        chunk_num = int(np.prod(batchedTensor.shape[1:-2]))
        if chunk_num >= (256*256 - 1):
            print("TOO BIG TENSOR")
        max_split = (256 * 256 - 1)//chunk_num
        det_list = []
        for t in torch.split(batchedTensor, max_split):
            det_list.append(torch.prod(torch.diagonal(torch.linalg.cholesky(t), dim1=-2, dim2=-1), dim = -1)**2)
        return torch.cat(det_list)
    else:
        return torch.prod(torch.diagonal(torch.linalg.cholesky(batchedTensor), dim1=-2, dim2=-1), dim = -1)**2
//  [python]
def test_batch_cg(self):
        batch = 5
        size = 100
        matrix = torch.randn(batch, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.transpose(-1, -2))
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(batch, size, 50, dtype=torch.float64)
        solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)

        # Check cg
        matrix_chol = torch.cholesky(matrix)
        actual = torch.cholesky_solve(rhs, matrix_chol)
        assert (torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
//  [python]
def test_batch_cg(self):
        batch = 5
        size = 100
        matrix = torch.randn(batch, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.transpose(-1, -2))
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(batch, size, 50, dtype=torch.float64)
        solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)

        # Check cg
        matrix_chol = torch.linalg.cholesky(matrix)
        actual = torch.linalg.solve_triangular(rhs, matrix_chol, upper=False)
        assert (torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))