jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _apply_floor_(self, A, B, component_batches):
        L = torch.linalg.cholesky(B)
        L_inv = torch.linalg.inv(L)
        num_floored = 0
        batch_size = component_batches[0][1] - component_batches[0][0]
        l = torch.zeros(batch_size, self.feat_dim, dtype=self.means.dtype, device=self.t_matrix.device)
        U = torch.zeros(batch_size, self.feat_dim, self.feat_dim, dtype=self.means.dtype, device=self.t_matrix.device) 
        for bstart, bend in component_batches:
            D = L_inv.matmul(A[bstart:bend, :, :]).matmul(L_inv.t())   
            for c in range(batch_size):
                l[c, :], U[c, :, :] = torch.symeig(D[c, :, :], eigenvectors=True)
            num_floored += torch.sum(l < 1).item()
            l = torch.clamp(l, min=1)
            D = U.matmul(l.unsqueeze(2) * U.transpose(1,2))
            A[bstart:bend, :, :] = L.matmul(D.transpose(1, 2)).matmul(L.t())
        print('Floored {:.1%} of the eigenvalues...'.format(num_floored / (self.n_components * self.feat_dim)))
//  [python]
def _distance(self, X, Y):
        L_inv = torch.inverse(torch.cholesky(X))
        return matrix_log(L_inv.dot(Y).dot(L_inv.T))
//  [python]
def _distance(self, X, Y):
        L_inv = torch.inverse(torch.linalg.cholesky(X))
        return matrix_log(L_inv.dot(Y).dot(L_inv.T))
//  [python]
def sigma(self, val):
        assert len(val.shape) == 3
        self._L = torch.cholesky(val)
//  [python]
def sigma(self, val):
        assert len(val.shape) == 3
        self._L = torch.linalg.cholesky(val)
//  [python]
def negELBO(self):
        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_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))
        K1 = torch.cat((torch.cat((K_u_u, K_u_tt.T), 1), torch.cat((K_u_tt, K_u_tt_u_tt), 1)), 0)
        K_gg = self.kernel.matrix(self.X_col, torch.exp(self.log_ls_h))
        K = torch.block_diag(K1, K_gg)
        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) - f[self.N_con + self.N_col:self.N_con + 2 * self.N_col].view(-1, 1))))
        return -elbo
//  [python]
def negELBO(self):
        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_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))
        K1 = torch.cat((torch.cat((K_u_u, K_u_tt.T), 1), torch.cat((K_u_tt, K_u_tt_u_tt), 1)), 0)
        K_gg = self.kernel.matrix(self.X_col, torch.exp(self.log_ls_h))
        K = torch.block_diag(K1, K_gg)
        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) - f[self.N_con + self.N_col:self.N_con + 2 * self.N_col].view(-1, 1))))
        return -elbo
//  [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(Xte)
        X = self.dkl_nn(self.X)

        Sigma = self.kernel(X, X) + self.log_beta.exp().pow(-1) * torch.eye(X.size(0)) \
            + 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(Xte)
        X = self.dkl_nn(self.X)

        Sigma = self.kernel(X, X) + self.log_beta.exp().pow(-1) * torch.eye(X.size(0)) \
            + 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 train(self, X, y, method='cholesky', alpha=1e-2):
        """
        Compute the output weights with a linear regression.
        Parameters:
          - X: input sequence of shape (seq_len, res_size)
          - y: target output (seq_len, out_dim)
          - method: "cholesky" or "sklearn ridge"
          - alpha: L2-regularization parameter
        
        Returns: a tensor of shape (res_size, out_dim)
        """
        if method == 'cholesky':
            # This technique uses the Cholesky decomposition
            # It should be fast when res_size < seq_len
            Xt_y = X.T @ y  # size (res_size, out_dim)
            K = X.T @ X  # size (res_size, res_size)
            K.view(-1)[::len(K)+1] += alpha  # add elements on the diagonal inplace
            L = torch.cholesky(K, upper=False)
            return torch.cholesky_solve(Xt_y, L, upper=False)
        elif method == 'sklearn ridge':
            from sklearn.linear_model import Ridge
            clf = Ridge(fit_intercept=False, alpha=alpha)
            clf.fit(X.cpu().numpy(), y.cpu().numpy())
            return torch.from_numpy(clf.coef_.T).to(self.device)
//  [python]
def train(self, X, y, method='cholesky', alpha=1e-2):
        """
        Compute the output weights with a linear regression.
        Parameters:
          - X: input sequence of shape (seq_len, res_size)
          - y: target output (seq_len, out_dim)
          - method: "cholesky" or "sklearn ridge"
          - alpha: L2-regularization parameter
        
        Returns: a tensor of shape (res_size, out_dim)
        """
        if method == 'cholesky':
            # This technique uses the Cholesky decomposition
            # It should be fast when res_size < seq_len
            Xt_y = X.T @ y  # size (res_size, out_dim)
            K = X.T @ X  # size (res_size, res_size)
            K.view(-1)[::len(K)+1] += alpha  # add elements on the diagonal inplace
            L = torch.linalg.cholesky(K, upper=False)
            return torch.linalg.solve_triangular(L, Xt_y, upper=False)
        elif method == 'sklearn ridge':
            from sklearn.linear_model import Ridge
            clf = Ridge(fit_intercept=False, alpha=alpha)
            clf.fit(X.cpu().numpy(), y.cpu().numpy())
            return torch.from_numpy(clf.coef_.T).to(self.device)
//  [python]
def sample(self, batchSize):
        L = torch.cholesky(self.sigma)
        size = [batchSize] + self.nvars
        return (self.mu+torch.matmul(torch.randn(size).to(self.sigma),self.sigma)).to(self.sigma)
//  [python]
def sample(self, batchSize):
        L = torch.linalg.cholesky(self.sigma)
        size = [batchSize] + self.nvars
        return (self.mu+torch.matmul(torch.randn(size).to(self.sigma),self.sigma)).to(self.sigma)
//  [python]
def plot_so3_helper(T_est, P_est, color, i1, i2, i_max, path):
    r = torch.randn(i_max, 9)
    r = r[r.norm(dim=1) < 4.1682][:, 6:9]
    P_est_chol = torch.cholesky(P_est[6:9, 6:9])
    p_est = bmv(P_est_chol.expand(r.shape[0], 3, 3), r) + T_est[:3, 4]
    plt.scatter(p_est[:, i1], p_est[:, i2], color=color, s=2, alpha=0.5)
    if i1 == 0 and i2 == 1:
        np.savetxt(path2, p_est.numpy(), header="x y z", comments='')
//  [python]
def plot_so3_helper(T_est, P_est, color, i1, i2, i_max, path):
    r = torch.randn(i_max, 9)
    r = r[r.norm(dim=1) < 4.1682][:, 6:9]
    P_est_chol = torch.linalg.cholesky(P_est[6:9, 6:9])
    p_est = bmv(P_est_chol.expand(r.shape[0], 3, 3), r) + T_est[:3, 4]
    plt.scatter(p_est[:, i1], p_est[:, i2], color=color, s=2, alpha=0.5)
    if i1 == 0 and i2 == 1:
        np.savetxt(path2, p_est.numpy(), header="x y z", comments='')
//  [python]
def pred(self, Xte):
        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_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))
        K1 = torch.cat((torch.cat((K_u_u, K_u_tt.T), 1), torch.cat((K_u_tt, K_u_tt_u_tt), 1)), 0)
        K_gg = self.kernel.matrix(self.X_col, torch.exp(self.log_ls_h))
        K = torch.block_diag(K1, K_gg)
        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(K1, f[0:self.N_con + self.N_col]))
        A1 = torch.cholesky(K1)
        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(K1, K_te.T).T).sum(1) + (K_te * torch.matmul(t2, K_te.T).T).sum(1))
        pred_stds = pred_stds.view(pred_mean.shape)
        return pred_mean, pred_stds**0.5
//  [python]
def pred(self, Xte):
        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_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))
        K1 = torch.cat((torch.cat((K_u_u, K_u_tt.T), 1), torch.cat((K_u_tt, K_u_tt_u_tt), 1)), 0)
        K_gg = self.kernel.matrix(self.X_col, torch.exp(self.log_ls_h))
        K = torch.block_diag(K1, K_gg)
        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(K1, f[0:self.N_con + self.N_col]))
        A1 = torch.linalg.cholesky(K1)
        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(K1, K_te.T).T).sum(1) + (K_te * torch.matmul(t2, K_te.T).T).sum(1))
        pred_stds = pred_stds.view(pred_mean.shape)
        return pred_mean, pred_stds**0.5
//  [python]
def estimate_A_prox_pytorch(n, n_aug, Y_h, Yh, rho, h, d,  return_P_inv, max_it, verbose,return_res=0):
    bs, _, _ = Y_h.shape
    A_t_s = []
    # Initialize A with identity, B with zeros
    A0 = torch.zeros((bs, n, n_aug), dtype=Yh.dtype, device=Yh.device)
    for i in range(h):
        A0[:,:, i * n:(i + 1) * n] = batch_eye(n=n, bs=bs, dtype=Yh.dtype, device=Yh.device)
    if d:
        A0[:, :,  -d:] = torch.zeros((bs, n, d), dtype=Yh.dtype, device=Yh.device)  # Initialize B component

    A_t_s.append(A0.transpose(1, 2))
    rho = rho

    P_b = Y_h@(Y_h.transpose(1, 2)) + rho * batch_eye(n=n_aug, bs=bs, dtype=Yh.dtype, device=Yh.device)
    L_b = torch.cholesky(P_b)
    if return_P_inv:
        P_inv = torch.cholesky_solve(batch_eye(n=h*n, bs=bs, dtype=Yh.dtype, device=Yh.device), L_b)
    for k in range(max_it):
        res = torch.norm(A_t_s[-1].transpose(1, 2) @ Y_h - Yh, np.inf)
        A_ = A_t_s[-1] + torch.cholesky_solve(Y_h.bmm(Yh.transpose(1, 2) - Y_h.transpose(1, 2).bmm(A_t_s[-1])), L_b)
        A_t_s.append(A_)
    A = A_t_s[-1].transpose(1, 2)
    if return_res:
        return A, res
    else:
        return A
//  [python]
def estimate_A_prox_pytorch(n, n_aug, Y_h, Yh, rho, h, d,  return_P_inv, max_it, verbose,return_res=0):
    bs, _, _ = Y_h.shape
    A_t_s = []
    # Initialize A with identity, B with zeros
    A0 = torch.zeros((bs, n, n_aug), dtype=Yh.dtype, device=Yh.device)
    for i in range(h):
        A0[:,:, i * n:(i + 1) * n] = batch_eye(n=n, bs=bs, dtype=Yh.dtype, device=Yh.device)
    if d:
        A0[:, :,  -d:] = torch.zeros((bs, n, d), dtype=Yh.dtype, device=Yh.device)  # Initialize B component

    A_t_s.append(A0.transpose(1, 2))
    rho = rho

    P_b = Y_h@(Y_h.transpose(1, 2)) + rho * batch_eye(n=n_aug, bs=bs, dtype=Yh.dtype, device=Yh.device)
    L_b = torch.linalg.cholesky(P_b)
    if return_P_inv:
        P_inv = torch.linalg.solve(torch.linalg.cholesky(P_b), batch_eye(n=h*n, bs=bs, dtype=Yh.dtype, device=Yh.device))
    for k in range(max_it):
        res = torch.norm(A_t_s[-1].transpose(1, 2) @ Y_h - Yh, np.inf)
        A_ = A_t_s[-1] + torch.linalg.solve(L_b, torch.linalg.solve(L_b.transpose(-2,-1), Y_h.bmm(Yh.transpose(1, 2) - Y_h.transpose(1, 2).bmm(A_t_s[-1]))))
        A_t_s.append(A_)
    A = A_t_s[-1].transpose(1, 2)
    if return_res:
        return A, res
    else:
        return A
//  [python]
def toMahalanobisDistance(cls, target, mean, pred_cov, clamp_covariance=False):
        """ Generic function that can be used once vec2Cov is implemented
        can be reimplemented if a better way to do it exists with one parametrization
        Args:
            mean [n x 3] : vx, vy, vz
            pred_cov [n x params] : xx, yy, zz
        Returns:
            err [n x 1] : mahalanobis distance square
        """
        cov_matrix = cls.vec2Cov(pred_cov)
        # compute the inverse of covariance matrices
        CovInv = torch.zeros_like(cov_matrix)
        N = target.shape[0]
        for i in range(N):
            u = torch.cholesky(cov_matrix[i, :, :])
            CovInv[i, :, :] = torch.cholesky_inverse(u)

        # compute the error
        err = mean - target
        loss_part1 = torch.einsum("ki,kij,kj->k", err, CovInv, err)
        if clamp_covariance:
            loss_part2 = torch.log(cov_matrix.det().clamp(min=1e-10))
        else:
            loss_part2 = torch.logdet(cov_matrix)

        loss = loss_part1 + loss_part2
        return loss.reshape((N, -1))
//  [python]
def toMahalanobisDistance(cls, target, mean, pred_cov, clamp_covariance=False):
        """ Generic function that can be used once vec2Cov is implemented
        can be reimplemented if a better way to do it exists with one parametrization
        Args:
            mean [n x 3] : vx, vy, vz
            pred_cov [n x params] : xx, yy, zz
        Returns:
            err [n x 1] : mahalanobis distance square
        """
        cov_matrix = cls.vec2Cov(pred_cov)
        # compute the inverse of covariance matrices
        CovInv = torch.zeros_like(cov_matrix)
        N = target.shape[0]
        for i in range(N):
            u = torch.linalg.cholesky(cov_matrix[i, :, :])
            CovInv[i, :, :] = torch.linalg.cholesky_inverse(u)

        # compute the error
        err = mean - target
        loss_part1 = torch.einsum("ki,kij,kj->k", err, CovInv, err)
        if clamp_covariance:
            loss_part2 = torch.log(cov_matrix.det().clamp(min=1e-10))
        else:
            loss_part2 = torch.logdet(cov_matrix)

        loss = loss_part1 + loss_part2
        return loss.reshape((N, -1))
//  [python]
def forward(ctx, A, b):
        # A: (..., M, N)
        # b: (..., M, K)
        # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/linalg_ops.py#L267
        u = torch.cholesky(torch.matmul(A.transpose(-1, -2), A), upper=True)
        ret = torch.cholesky_solve(torch.matmul(A.transpose(-1, -2), b), u, upper=True)
        ctx.save_for_backward(u, ret, A, b)
        return ret
//  [python]
def forward(ctx, A, b):
        # A: (..., M, N)
        # b: (..., M, K)
        # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/linalg_ops.py#L267
        u = torch.linalg.cholesky(torch.matmul(A.transpose(-1, -2), A))
        ret = torch.linalg.cholesky_solve(torch.matmul(A.transpose(-1, -2), b), u)
        ctx.save_for_backward(u, ret, A, b)
        return ret
//  [python]
def test_cholupdate_negative():
    """Checks rank-1 Cholesky update forward pass, with weight set to -2."""
    batch_dims = tuple()
    matrix_dim = 3
    L = fannypack.utils.tril_from_vector(
        torch.randn(
            batch_dims + (fannypack.utils.tril_count_from_matrix_dim(matrix_dim),)
        )
    ) + torch.eye(matrix_dim)
    x = torch.ones(matrix_dim) * 0.2

    # Make sure our output will be PSD
    L = torch.cholesky(L @ L.transpose(-1, -2) + x[..., :, None] @ x[..., None, :])

    updated_L = fannypack.utils.cholupdate(
        L, x, weight=torch.tensor(-0.5, dtype=torch.float32)
    )

    torch.testing.assert_allclose(
        L @ L.transpose(-1, -2) - 0.5 * x[..., :, None] @ x[..., None, :],
        updated_L @ updated_L.transpose(-1, -2),
    )
//  [python]
def test_cholupdate_negative():
    """Checks rank-1 Cholesky update forward pass, with weight set to -2."""
    batch_dims = tuple()
    matrix_dim = 3
    L = fannypack.utils.tril_from_vector(
        torch.randn(
            batch_dims + (fannypack.utils.tril_count_from_matrix_dim(matrix_dim),)
        )
    ) + torch.eye(matrix_dim)
    x = torch.ones(matrix_dim) * 0.2

    # Make sure our output will be PSD
    L = torch.linalg.cholesky(L @ L.transpose(-1, -2) + x[..., :, None] @ x[..., None, :])

    updated_L = fannypack.utils.cholupdate(
        L, x, weight=torch.tensor(-0.5, dtype=torch.float32)
    )

    torch.testing.assert_allclose(
        L @ L.transpose(-1, -2) - 0.5 * x[..., :, None] @ x[..., None, :],
        updated_L @ updated_L.transpose(-1, -2),
    )
//  [python]
def test_cholupdate_negative_float_weight():
    """Checks rank-1 Cholesky update forward pass, with weight set to -2."""
    batch_dims = tuple()
    matrix_dim = 3
    L = fannypack.utils.tril_from_vector(
        torch.randn(
            batch_dims + (fannypack.utils.tril_count_from_matrix_dim(matrix_dim),)
        )
    ) + torch.eye(matrix_dim)
    x = torch.ones(matrix_dim) * 0.2

    # Make sure our output will be PSD
    L = torch.cholesky(L @ L.transpose(-1, -2) + x[..., :, None] @ x[..., None, :])

    updated_L = fannypack.utils.cholupdate(L, x, weight=-0.5)

    torch.testing.assert_allclose(
        L @ L.transpose(-1, -2) - 0.5 * x[..., :, None] @ x[..., None, :],
        updated_L @ updated_L.transpose(-1, -2),
    )
//  [python]
def test_cholupdate_negative_float_weight():
    """Checks rank-1 Cholesky update forward pass, with weight set to -2."""
    batch_dims = tuple()
    matrix_dim = 3
    L = fannypack.utils.tril_from_vector(
        torch.randn(
            batch_dims + (fannypack.utils.tril_count_from_matrix_dim(matrix_dim),)
        )
    ) + torch.eye(matrix_dim)
    x = torch.ones(matrix_dim) * 0.2

    # Make sure our output will be PSD
    L = torch.linalg.cholesky(L @ L.transpose(-1, -2) + x[..., :, None] @ x[..., None, :])

    updated_L = fannypack.utils.cholupdate(L, x, weight=-0.5)

    torch.testing.assert_allclose(
        L @ L.transpose(-1, -2) - 0.5 * x[..., :, None] @ x[..., None, :],
        updated_L @ updated_L.transpose(-1, -2),
    )
//  [python]
def _batch_conditioned_dynamics(nx, nu, sigma, mu, sigreg=1e-5):
    # TODO assume we want to calculate everything to the right of xu (can't handle pxu)
    it = slice(nx + nu)
    ip = slice(nx + nu, None)
    # guarantee symmetric positive definite with regularization
    sigma[:, it, it] += sigreg * torch.eye(nx + nu, dtype=sigma.dtype, device=sigma.device)

    # sometimes fails for no good reason on the GPU; try a couple of times
    u = None
    for t in range(3):
        try:
            u = torch.cholesky(sigma[:, it, it])
            # if we succeeded just exit for loop
            break
        except RuntimeError as e:
            logger.warning(e)

    if u is None:
        raise RuntimeError("Failed cholesky multiple times; actual problem with data")

    # equivalent to inv * sigma
    # Solve normal equations to get dynamics. (equation 2)
    Fm = torch.cholesky_solve(sigma[:, it, ip], u).transpose(-1, -2)
    fv = mu[:, ip] - linalg.batch_batch_product(mu[:, it], Fm)

    # TODO calculate dyn_covar
    return Fm, fv, None
//  [python]
def _batch_conditioned_dynamics(nx, nu, sigma, mu, sigreg=1e-5):
    # TODO assume we want to calculate everything to the right of xu (can't handle pxu)
    it = slice(nx + nu)
    ip = slice(nx + nu, None)
    # guarantee symmetric positive definite with regularization
    sigma[:, it, it] += sigreg * torch.eye(nx + nu, dtype=sigma.dtype, device=sigma.device)

    # sometimes fails for no good reason on the GPU; try a couple of times
    u = None
    for t in range(3):
        try:
            u = torch.linalg.cholesky(sigma[:, it, it])
            # if we succeeded just exit for loop
            break
        except RuntimeError as e:
            logger.warning(e)

    if u is None:
        raise RuntimeError("Failed cholesky multiple times; actual problem with data")

    # equivalent to inv * sigma
    # Solve normal equations to get dynamics. (equation 2)
    Fm = torch.linalg.solve_triangular(u, sigma[:, it, ip], upper=False).transpose(-1, -2)
    fv = mu[:, ip] - linalg.batch_batch_product(mu[:, it], Fm)

    # TODO calculate dyn_covar
    return Fm, fv, None
//  [python]
def pred(self, Xte):
        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))
        K1 = torch.cat((torch.cat((K_u_u, K_utt.T), 1), torch.cat((K_utt, K_utt_utt), 1)), 0)
        K_gg = self.kernel.matrix(self.X_col, torch.exp(self.log_ls_h))
        K = torch.block_diag(K1, K_gg)
        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(K1, f[0:self.N_con + self.N_col]))
        A1 = torch.cholesky(K1)
        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(K1, K_te.T).T).sum(1) + (K_te * torch.matmul(t2, K_te.T).T).sum(1))
        pred_stds = pred_stds.view(pred_mean.shape)
        return pred_mean, pred_stds**0.5
//  [python]
def pred(self, Xte):
        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))
        K1 = torch.cat((torch.cat((K_u_u, K_utt.T), 1), torch.cat((K_utt, K_utt_utt), 1)), 0)
        K_gg = self.kernel.matrix(self.X_col, torch.exp(self.log_ls_h))
        K = torch.block_diag(K1, K_gg)
        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(K1, f[0:self.N_con + self.N_col]))
        A1 = torch.linalg.cholesky(K1)
        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(K1, K_te.T).T).sum(1) + (K_te * torch.matmul(t2, K_te.T).T).sum(1))
        pred_stds = pred_stds.view(pred_mean.shape)
        return pred_mean, pred_stds**0.5
//  [python]
def _fit_data(self):
        # Fit data
        self.K = self.kernel(self.train_x, self.train_x) + \
                 self.noise_variance() * th.eye(self.n_samples,
                                                device=self.train_x.device)
        self.L = th.cholesky(self.K)
        self.alpha = th.cholesky_solve(self.train_y, self.L)
//  [python]
def _fit_data(self):
        # Fit data
        self.K = self.kernel(self.train_x, self.train_x) + \
                 self.noise_variance() * th.eye(self.n_samples,
                                                device=self.train_x.device)
        self.L = th.linalg.cholesky(self.K)
        self.alpha = th.linalg.solve_triangular(self.train_y, self.L, upper=False)
//  [python]
def forward(self, x_test, Kinv_method='cholesky3'):
        """
        Computes the mean and covariance of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            x_test (torch.Tensor): The input test data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

        Returns:
            tuple: A tuple containing the mean and covariance of the Gaussian process.

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        x_train = self.train_x
        y_train = self.train_y
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        K_s = self.kernel(x_train, x_test)
        K_ss = self.kernel(x_test, x_test)
        
        if Kinv_method == 'cholesky1':
            # kernel inverse is not stable, use cholesky decomposition instead
            L = torch.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            alpha = K_inv @ y_train
            mu = K_s.T @ alpha
            v = L_inv @ K_s
            cov = K_ss - v.T @ v
        elif Kinv_method == 'cholesky3':
            # recommended implementation, fastest so far
            L = torch.cholesky(K)
            alpha = torch.cholesky_solve(y_train, L)
            mu = K_s.T @ alpha
            v = L.inverse() @ K_s
            cov = K_ss - v.T @ v
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            mu = K_s.T @ K_inv @ y_train
            cov = K_ss - K_s.T @ K_inv @ K_s
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
        return mu.squeeze(), cov
//  [python]
def forward(self, x_test, Kinv_method='cholesky3'):
        """
        Computes the mean and covariance of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            x_test (torch.Tensor): The input test data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

        Returns:
            tuple: A tuple containing the mean and covariance of the Gaussian process.

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        x_train = self.train_x
        y_train = self.train_y
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        K_s = self.kernel(x_train, x_test)
        K_ss = self.kernel(x_test, x_test)
        
        if Kinv_method == 'cholesky1':
            # kernel inverse is not stable, use cholesky decomposition instead
            L = torch.linalg.cholesky(K)
            L_inv = torch.linalg.inv(L)
            K_inv = L_inv.T @ L_inv
            alpha = K_inv @ y_train
            mu = K_s.T @ alpha
            v = L_inv @ K_s
            cov = K_ss - v.T @ v
        elif Kinv_method == 'cholesky3':
            # recommended implementation, fastest so far
            L = torch.linalg.cholesky(K)
            alpha = torch.linalg.solve(L, y_train)
            mu = K_s.T @ alpha
            v = L.inverse() @ K_s
            cov = K_ss - v.T @ v
        elif Kinv_method == 'direct':
            K_inv = torch.linalg.inv(K)
            mu = K_s.T @ K_inv @ y_train
            cov = K_ss - K_s.T @ K_inv @ K_s
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
        return mu.squeeze(), cov
//  [python]
def test_cholesky(self):
        def func(root, upper):
            x = torch.matmul(root, root.transpose(-1, -2)) + 1e-05
            return torch.cholesky(x, upper)

        def run_test(upper, dims):
            root = torch.rand(*dims, requires_grad=True)

            gradcheck(func, [root, upper])
            gradgradcheck(func, [root, upper])

            root = random_symmetric_pd_matrix(dims[-1], *dims[:-2]).requires_grad_()
            chol = root.cholesky().sum().backward()
            self.assertEqual(root.grad, root.grad.transpose(-1, -2))  # Check the gradient is symmetric

        for upper, dims in product([True, False], [(3, 3), (4, 3, 2, 2)]):
            run_test(upper, dims)
            run_test(upper, dims)
//  [python]
def test_cholesky(self):
        def func(root, upper):
            x = torch.matmul(root, root.transpose(-1, -2)) + 1e-05
            return torch.linalg.cholesky(x, upper=upper)

        def run_test(upper, dims):
            root = torch.rand(*dims, requires_grad=True)

            gradcheck(func, [root, upper])
            gradgradcheck(func, [root, upper])

            root = random_symmetric_pd_matrix(dims[-1], *dims[:-2]).requires_grad_()
            chol = root.cholesky().sum().backward()
            self.assertEqual(root.grad, root.grad.transpose(-1, -2))  # Check the gradient is symmetric

        for upper, dims in product([True, False], [(3, 3), (4, 3, 2, 2)]):
            run_test(upper, dims)
            run_test(upper, dims)
//  [python]
def _posdef_logdet(A: torch.Tensor) -> torch.Tensor:
    tril = torch.cholesky(A, upper=False)
    return 2. * _triangular_logdet(tril)
//  [python]
def _posdef_logdet(A: torch.Tensor) -> torch.Tensor:
    tril = torch.linalg.cholesky(A, upper=False)
    return 2. * _triangular_logdet(tril)
//  [python]
def scale_tril(self):
        return torch.cholesky(self.scale, upper=False)
//  [python]
def scale_tril(self):
        return torch.linalg.cholesky(self.scale, upper=False)
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
    """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)
    """
    try:
        L = torch.cholesky(A, upper=upper, out=out)
        return L
    except RuntimeError as e:
        isnan = torch.isnan(A)
        if isnan.any():
            raise NanError(
                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(5):
            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} to the diagonal",
                    RuntimeWarning,
                )
                return L
            except RuntimeError:
                continue
        raise e
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
    """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)
    """
    try:
        L = torch.linalg.cholesky(A, upper=upper, out=out)
        return L
    except RuntimeError as e:
        isnan = torch.isnan(A)
        if isnan.any():
            raise NanError(
                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(5):
            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} to the diagonal",
                    RuntimeWarning,
                )
                return L
            except RuntimeError:
                continue
        raise e
//  [python]
def log_likelihood(self, x_train, y_train, Kinv_method='cholesky3'):
        """
        Computes the log-likelihood of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

        Returns:
            torch.Tensor: The log-likelihood of the Gaussian process.

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        self.train_x = x_train
        self.train_y = y_train

        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        
        if Kinv_method == 'cholesky1':
            L = torch.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            return -0.5 * (y_train.T @ K_inv @ y_train + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky2':
            L = torch.cholesky(K)
            return -0.5 * (y_train.T @ torch.cholesky_solve(y_train, L) + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky3':
            L = torch.cholesky(K)
            return -0.5 * (y_train.T @ torch.cholesky_solve(y_train, L) + L.diag().log().sum() + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            return -0.5 * (y_train.T @ K_inv @ y_train + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'torch_distribution_MN1':
            L = torch.cholesky(K)
            return torch.distributions.MultivariateNormal(y_train, scale_tril=L).log_prob(y_train)
        elif Kinv_method == 'torch_distribution_MN2':
            return torch.distributions.MultivariateNormal(y_train, K).log_prob(y_train)
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
//  [python]
def log_likelihood(self, x_train, y_train, Kinv_method='cholesky3'):
        """
        Computes the log-likelihood of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

        Returns:
            torch.Tensor: The log-likelihood of the Gaussian process.

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        self.train_x = x_train
        self.train_y = y_train

        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        
        if Kinv_method == 'cholesky1':
            L = torch.linalg.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            return -0.5 * (y_train.T @ K_inv @ y_train + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky2':
            L = torch.linalg.cholesky(K)
            return -0.5 * (y_train.T @ torch.cholesky_solve(y_train, L) + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky3':
            L = torch.linalg.cholesky(K)
            return -0.5 * (y_train.T @ torch.cholesky_solve(y_train, L) + L.diag().log().sum() + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            return -0.5 * (y_train.T @ K_inv @ y_train + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'torch_distribution_MN1':
            L = torch.linalg.cholesky(K)
            return torch.distributions.MultivariateNormal(y_train, scale_tril=L).log_prob(y_train)
        elif Kinv_method == 'torch_distribution_MN2':
            return torch.distributions.MultivariateNormal(y_train, K).log_prob(y_train)
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
//  [python]
def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        D = self.event_shape[-1]
        value_tril = torch.cholesky(value, upper=False)
        logdet = 2. * _triangular_logdet(value_tril)
        # TODO: Replace with cholesky_solve once its derivative is implemented
        # prod = torch.cholesky_solve(self.scale, value_tril, upper=False)
        prod_ = torch.triangular_solve(self.scale, value_tril, upper=False)[0]
        prod = torch.triangular_solve(prod_, value_tril, upper=False, transpose=True)[0]
        trace = prod.diagonal(dim1=-2, dim2=-1).sum(-1)
        return -.5 * (self.df + D + 1) * logdet - .5 * trace - self.log_normalizer
//  [python]
def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        D = self.event_shape[-1]
        value_tril = torch.linalg.cholesky(value, upper=False)
        logdet = 2. * _triangular_logdet(value_tril)
        # TODO: Replace with cholesky_solve once its derivative is implemented
        # prod = torch.cholesky_solve(self.scale, value_tril, upper=False)
        prod_ = torch.triangular_solve(self.scale, value_tril, upper=False)[0]
        prod = torch.triangular_solve(prod_, value_tril, upper=False, transpose=True)[0]
        trace = prod.diagonal(dim1=-2, dim2=-1).sum(-1)
        return -.5 * (self.df + D + 1) * logdet - .5 * trace - self.log_normalizer
//  [python]
def step(self, update_stats=True, update_params=True):
        #Performs one step of preconditioning.
        fisher_norm = 0.
        for group in self.param_groups:
            #print(torch.cuda.memory_allocated()/(1024**3))
            # Getting parameters
            if len(group['params']) == 2:
                weight, bias = group['params']
            else:
                weight = group['params'][0]
                bias = None
            state = self.state[weight]

            #print(group['layer_type'])
            if 'BatchNorm' in group['layer_type'] and self.use_batch_norm:
                # now compute hessian of weights
                diag_comp = 100 * weight.size(0) * self.eps * torch.eye(weight.size(0), device = weight.device, dtype = weight.dtype)
                #print(weight.size())
                weight_hessian = jacobian(weight.grad, weight) + diag_comp
                #print(weight_hessian)

                weight_inv_chol = torch.cholesky(weight_hessian)
                state['w_ic'] = weight_inv_chol

                if bias is not None:
                    diag_comp = 100 * self.eps * torch.eye(bias.size(0), device = bias.device, dtype = bias.dtype)
                    bias_hessian = jacobian(bias.grad, bias) + diag_comp

                    state['b_ic'] = torch.cholesky(bias_hessian)

            if group['layer_type'] in ['Linear', 'Conv2d']:

                # Update convariances and inverses
                if update_stats:
                    if self._iteration_counter % self.update_freq == 0:
                        self._compute_covs(group, state)
                        ixxt, iggt, _, _ = self._inv_covs(state['xxt'], state['ggt'],
                                                    state['num_locations'])
                        state['ixxt'] = ixxt
                        state['iggt'] = iggt
                    else:
                        if self.alpha != 1:
                            self._compute_covs(group, state)
                if update_params:
                    # Preconditionning
                    gw, gb = self._precond(weight, bias, group, state)
                    # Updating gradients
                    if self.constraint_norm:
                        fisher_norm += (weight.grad * gw).sum()
                    weight.grad.data = gw
                    if bias is not None:
                        if self.constraint_norm:
                            fisher_norm += (bias.grad * gb).sum()
                        bias.grad.data = gb
                # Cleaning
                if 'x' in self.state[group['mod']]:
                    del self.state[group['mod']]['x']
                if 'gy' in self.state[group['mod']]:
                    del self.state[group['mod']]['gy']
        # Eventually scale the norm of the gradients
        if update_params and self.constraint_norm:
            f_scale = (1. / fisher_norm) ** 0.5
            for group in self.param_groups:
                for param in group['params']:
                    param.grad.data *= f_scale
        if update_stats:
            self._iteration_counter += 1
//  [python]
def step(self, update_stats=True, update_params=True):
        #Performs one step of preconditioning.
        fisher_norm = 0.
        for group in self.param_groups:
            #print(torch.cuda.memory_allocated()/(1024**3))
            # Getting parameters
            if len(group['params']) == 2:
                weight, bias = group['params']
            else:
                weight = group['params'][0]
                bias = None
            state = self.state[weight]

            #print(group['layer_type'])
            if 'BatchNorm' in group['layer_type'] and self.use_batch_norm:
                # now compute hessian of weights
                diag_comp = 100 * weight.size(0) * self.eps * torch.eye(weight.size(0), device = weight.device, dtype = weight.dtype)
                #print(weight.size())
                weight_hessian = jacobian(weight.grad, weight) + diag_comp
                #print(weight_hessian)

                weight_inv_chol = torch.linalg.cholesky(weight_hessian)
                state['w_ic'] = weight_inv_chol

                if bias is not None:
                    diag_comp = 100 * self.eps * torch.eye(bias.size(0), device = bias.device, dtype = bias.dtype)
                    bias_hessian = jacobian(bias.grad, bias) + diag_comp

                    state['b_ic'] = torch.linalg.cholesky(bias_hessian)

            if group['layer_type'] in ['Linear', 'Conv2d']:

                # Update convariances and inverses
                if update_stats:
                    if self._iteration_counter % self.update_freq == 0:
                        self._compute_covs(group, state)
                        ixxt, iggt, _, _ = self._inv_covs(state['xxt'], state['ggt'],
                                                    state['num_locations'])
                        state['ixxt'] = ixxt
                        state['iggt'] = iggt
                    else:
                        if self.alpha != 1:
                            self._compute_covs(group, state)
                if update_params:
                    # Preconditionning
                    gw, gb = self._precond(weight, bias, group, state)
                    # Updating gradients
                    if self.constraint_norm:
                        fisher_norm += (weight.grad * gw).sum()
                    weight.grad.data = gw
                    if bias is not None:
                        if self.constraint_norm:
                            fisher_norm += (bias.grad * gb).sum()
                        bias.grad.data = gb
                # Cleaning
                if 'x' in self.state[group['mod']]:
                    del self.state[group['mod']]['x']
                if 'gy' in self.state[group['mod']]:
                    del self.state[group['mod']]['gy']
        # Eventually scale the norm of the gradients
        if update_params and self.constraint_norm:
            f_scale = (1. / fisher_norm) ** 0.5
            for group in self.param_groups:
                for param in group['params']:
                    param.grad.data *= f_scale
        if update_stats:
            self._iteration_counter += 1