jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def direct_newton(env, cmd, reg_ctrl, variant='gmres_scipy', time_solve_only=False):
    tic0 = time.time()
    horizon, dim_ctrl = cmd.shape

    # Compute step from a given command
    cmd_flat = deepcopy(cmd.data)
    cmd_flat = cmd_flat.view(-1)
    cmd_flat.requires_grad = True
    cmd_aux = cmd_flat.view(horizon, dim_ctrl)

    _, costs = env.forward(cmd_aux, approx='linquad')
    total_cost = sum(costs)

    # Get gradient, hessian, and make a newton step to get the solution
    grad = torch.autograd.grad(total_cost, cmd_flat, create_graph=True)[0]
    # add regularization in the hessian
    if variant == 'torch_solve':
      hess = auto_multi_grad(grad, cmd_flat) + reg_ctrl * torch.eye(dim_ctrl * horizon)
      tic1 = time.time()
      cmd_opt_newton = - torch.linalg.solve(hess, grad.unsqueeze(-1), ).view(-1)
    elif variant == 'own_cg':
      hvp = make_hvp(grad, cmd_flat)
      hvp_reg = lambda x: hvp(x) + reg_ctrl*x
      tic1 = time.time()
      cmd_opt_newton = own_cg(hvp_reg, grad)
    elif variant == 'gmres_scipy':
      hvp = make_numpy_hvp(grad, cmd_flat)
      hvp_reg = lambda x: hvp(x) + reg_ctrl*x
      hvp_reg = LinearOperator((horizon*dim_ctrl, horizon*dim_ctrl), hvp_reg)
      target = deepcopy(grad.data).numpy()
      tic1 = time.time()
      cmd_opt_newton = cg(hvp_reg, target)[0]
      cmd_opt_newton = torch.from_numpy(cmd_opt_newton)
    else:
        raise NotImplementedError


    cmd_opt_newton = cmd_opt_newton.view(horizon, dim_ctrl)
    time_newton = time.time() - tic1 if time_solve_only else time.time() - tic0
    return cmd_opt_newton, time_newton
//  [python]
def fit(self, X, Y):
        r"""
        Solve the Koopman Nystroem problem for KRR:

        .. math:: (K^x_{mn} K^x_{nm} + n\lambda K^x_{mm})^{-1} K^x_{mn} K^y_{nm} K^y_{mm}^{-1} f(Y)

        where for now :math:`f` is the identity function

        Parameters
        ----------
        X : torch.Tensor
            Covariates. These should be points in a time-series
        Y : torch.Tensor
            Targets. The corresponding future data in the time-series
        """
        ny_pts_x, ny_pts_y = self.center_selection.select(X, Y)
        self.fitX_, self.fitY_ = X, Y
        self.ny_pts_x_, self.ny_pts_y_ = ny_pts_x, ny_pts_y

        penalty = self.penalty * X.shape[0]

        kx_nm = self.kernel(X, ny_pts_x)
        kx_mm = self.kernel(ny_pts_x, ny_pts_x)

        ky_nm = self.kernel(Y, ny_pts_y)
        ky_mm = self.kernel(ny_pts_y, ny_pts_y)

        solve1 = torch.linalg.lstsq(ky_mm, ny_pts_y, driver='gels').solution

        rhs = kx_nm.T @ ky_nm @ solve1
        lhs = kx_nm.T @ kx_nm + penalty * kx_mm
        solve2 = torch.solve(rhs, lhs)[0]
        self.alpha_ = solve2
        return self
//  [python]
def fit(self, X, Y):
        r"""
        Solve the Koopman Nystroem problem for KRR:

        .. math:: (K^x_{mn} K^x_{nm} + n\lambda K^x_{mm})^{-1} K^x_{mn} K^y_{nm} K^y_{mm}^{-1} f(Y)

        where for now :math:`f` is the identity function

        Parameters
        ----------
        X : torch.Tensor
            Covariates. These should be points in a time-series
        Y : torch.Tensor
            Targets. The corresponding future data in the time-series
        """
        ny_pts_x, ny_pts_y = self.center_selection.select(X, Y)
        self.fitX_, self.fitY_ = X, Y
        self.ny_pts_x_, self.ny_pts_y_ = ny_pts_x, ny_pts_y

        penalty = self.penalty * X.shape[0]

        kx_nm = self.kernel(X, ny_pts_x)
        kx_mm = self.kernel(ny_pts_x, ny_pts_x)

        ky_nm = self.kernel(Y, ny_pts_y)
        ky_mm = self.kernel(ny_pts_y, ny_pts_y)

        solve1 = torch.linalg.lstsq(ky_mm, ny_pts_y, driver='gels').solution

        rhs = kx_nm.T @ ky_nm @ solve1
        lhs = kx_nm.T @ kx_nm + penalty * kx_mm
        solve2 = torch.linalg.solve(lhs, rhs)
        self.alpha_ = solve2
        return self
//  [python]
def ldet_sw_full(A, U, B, V):
    """
    Compute the log determinant of a matrix product:

    M = (A + UBV)

    Parameters
    ----------
    A : torch.Tensor
        Input matrix A.

    U : torch.Tensor
        Input matrix U.

    B : torch.Tensor
        Input matrix B.

    V : torch.Tensor
        Input matrix V.

    Returns
    -------
    torch.Tensor
        Log determinant of the matrix product.
    """
    # Compute the log determinant of matrix A
    s1, ldet_A = torch.slogdet(A)

    # Compute the log determinant of matrix B
    s2, ldet_B = torch.slogdet(B)

    # Compute the second term involving matrix V, A, and U
    term2 = torch.matmul(V, torch.solve(U, A)[0])

    # Compute the inverse of matrix B
    term1 = torch.inverse(B)

    # Compute the log determinant of the matrix product
    s3, ldet_prod = torch.slogdet(term1 + term2)

    # Compute the final log determinant
    log_determinant = ldet_A + ldet_B + ldet_prod

    return log_determinant
//  [python]
def ldet_sw_full(A, U, B, V):
    """
    Compute the log determinant of a matrix product:

    M = (A + UBV)

    Parameters
    ----------
    A : torch.Tensor
        Input matrix A.

    U : torch.Tensor
        Input matrix U.

    B : torch.Tensor
        Input matrix B.

    V : torch.Tensor
        Input matrix V.

    Returns
    -------
    torch.Tensor
        Log determinant of the matrix product.
    """
    # Compute the log determinant of matrix A
    s1, ldet_A = torch.slogdet(A)

    # Compute the log determinant of matrix B
    s2, ldet_B = torch.slogdet(B)

    # Compute the second term involving matrix V, A, and U
    term2 = torch.matmul(V, la.solve(A, U))

    # Compute the inverse of matrix B
    term1 = la.inv(B)

    # Compute the log determinant of the matrix product
    s3, ldet_prod = torch.slogdet(term1 + term2)

    # Compute the final log determinant
    log_determinant = ldet_A + ldet_B + ldet_prod

    return log_determinant
//  [python]
def test_invkernel(self):
        ############################################################
        import torch
        from pykeops.torch.operations import KernelSolve

        formula = "Exp(-oos2*SqDist(x,y))*b"
        aliases = [
            "x = Vi(" + str(self.D) + ")",  # First arg   : i-variable, of size D
            "y = Vj(" + str(self.D) + ")",  # Second arg  : j-variable, of size D
            "b = Vj(" + str(self.E) + ")",  # Third arg  : j-variable, of size Dv
            "oos2 = Pm(1)",
        ]  # Fourth arg  : scalar parameter

        Kinv = KernelSolve(formula, aliases, "b", axis=1)

        c = Kinv(self.xc, self.xc, self.ac, self.sigmac, alpha=self.alphac)
        torchsolve = lambda A, B: torch.solve(B, A)[0]
        c_ = torchsolve(
            self.alphac * torch.eye(self.M, device=self.device)
            + torch.exp(
                -torch.sum((self.xc[:, None, :] - self.xc[None, :, :]) ** 2, dim=2)
                * self.sigmac
            ),
            self.ac,
        )

        self.assertTrue(
            np.allclose(
                c.cpu().data.numpy().ravel(), c_.cpu().data.numpy().ravel(), atol=1e-4
            )
        )

        (u,) = torch.autograd.grad(c, self.xc, self.ec)
        (u_,) = torch.autograd.grad(c_, self.xc, self.ec)
        self.assertTrue(
            np.allclose(
                u.cpu().data.numpy().ravel(), u_.cpu().data.numpy().ravel(), atol=1e-4
            )
        )
//  [python]
def test_invkernel(self):
        ############################################################
        import torch
        from pykeops.torch.operations import KernelSolve

        formula = "Exp(-oos2*SqDist(x,y))*b"
        aliases = [
            "x = Vi(" + str(self.D) + ")",  # First arg   : i-variable, of size D
            "y = Vj(" + str(self.D) + ")",  # Second arg  : j-variable, of size D
            "b = Vj(" + str(self.E) + ")",  # Third arg  : j-variable, of size Dv
            "oos2 = Pm(1)",
        ]  # Fourth arg  : scalar parameter

        Kinv = KernelSolve(formula, aliases, "b", axis=1)

        c = Kinv(self.xc, self.xc, self.ac, self.sigmac, alpha=self.alphac)
        if torch.__version__ >= "1.8":
            torchsolve = lambda A, B: torch.linalg.solve(A, B)
        else:
            torchsolve = lambda A, B: torch.solve(B, A)[0]
        c_ = torchsolve(
            self.alphac * torch.eye(self.M, device=self.device)
            + torch.exp(
                -torch.sum((self.xc[:, None, :] - self.xc[None, :, :]) ** 2, dim=2)
                * self.sigmac
            ),
            self.ac,
        )

        self.assertTrue(
            np.allclose(
                c.cpu().data.numpy().ravel(), c_.cpu().data.numpy().ravel(), atol=1e-4
            )
        )

        (u,) = torch.autograd.grad(c, self.xc, self.ec)
        (u_,) = torch.autograd.grad(c_, self.xc, self.ec)
        self.assertTrue(
            np.allclose(
                u.cpu().data.numpy().ravel(), u_.cpu().data.numpy().ravel(), atol=1e-4
            )
        )
//  [python]
def _project(reference_sources, estimated_source, flen):
    """Least-squares projection of estimated source on the subspace spanned by
    delayed versions of reference sources, with delays between 0 and flen-1
    """
    nsrc = reference_sources.shape[0]
    nsampl = reference_sources.shape[1]
    kw = {'device': reference_sources.device, "dtype": reference_sources.dtype}
    # computing coefficients of least squares problem via FFT ##
    # zero padding and FFT of input data
    reference_sources = torch.cat((reference_sources,
                                   torch.zeros((nsrc, flen - 1), **kw)), dim=1)
    estimated_source = torch.cat((estimated_source, torch.zeros(flen - 1, **kw)), dim=0)
    n_fft = int(2 ** ceil(log2(nsampl + flen - 1.)))
    rs = _fix_shape(reference_sources, n_fft, -1)  # Padding like scipy.fftpack.fft does
    sf = fft(rs)
    es = _fix_shape(estimated_source, n_fft, -1)  # Padding like scipy.fftpack.fft does
    sef = fft(es)
    # inner products between delayed versions of reference_sources
    D = _calc_D(sf, sef, nsrc, flen, **kw)
    # inner products between estimated_source and delayed versions of
    # reference_sources
    G = _calc_G(sf, nsrc, flen, **kw)

    # Computing projection
    # Distortion filters
    if torch.det(G) > 0.1:
        C = torch.solve(D.unsqueeze(1), G).solution.reshape(nsrc, flen).T
    else:
        C = torch.lstsq(G,D.unsqueeze(1)).solution.reshape(nsrc, flen).T
    # Filtering
    sproj = torch.zeros(nsampl + flen - 1, **kw)
    for i in range(nsrc):
        sproj += torch.nn.functional.conv1d(reference_sources[i][:-flen + 1][None, None, ...],
                                            C[:, i].flip([0])[None, None, ...],
                                            padding=flen - 1)[0, 0]

    return sproj
//  [python]
def _project(reference_sources, estimated_source, flen):
    """Least-squares projection of estimated source on the subspace spanned by
    delayed versions of reference sources, with delays between 0 and flen-1
    """
    nsrc = reference_sources.shape[0]
    nsampl = reference_sources.shape[1]
    kw = {'device': reference_sources.device, "dtype": reference_sources.dtype}
    # computing coefficients of least squares problem via FFT ##
    # zero padding and FFT of input data
    reference_sources = torch.cat((reference_sources,
                                   torch.zeros((nsrc, flen - 1), **kw)), dim=1)
    estimated_source = torch.cat((estimated_source, torch.zeros(flen - 1, **kw)), dim=0)
    n_fft = int(2 ** ceil(log2(nsampl + flen - 1.)))
    rs = _fix_shape(reference_sources, n_fft, -1)  # Padding like scipy.fftpack.fft does
    sf = fft(rs)
    es = _fix_shape(estimated_source, n_fft, -1)  # Padding like scipy.fftpack.fft does
    sef = fft(es)
    # inner products between delayed versions of reference_sources
    D = _calc_D(sf, sef, nsrc, flen, **kw)
    # inner products between estimated_source and delayed versions of
    # reference_sources
    G = _calc_G(sf, nsrc, flen, **kw)

    # Computing projection
    # Distortion filters
    if torch.det(G) > 0.1:
        C = torch.linalg.solve(G, D.unsqueeze(1)).reshape(nsrc, flen).T
    else:
        C = torch.linalg.lstsq(G,D.unsqueeze(1)).solution.reshape(nsrc, flen).T
    # Filtering
    sproj = torch.zeros(nsampl + flen - 1, **kw)
    for i in range(nsrc):
        sproj += torch.nn.functional.conv1d(reference_sources[i][:-flen + 1][None, None, ...],
                                            C[:, i].flip([0])[None, None, ...],
                                            padding=flen - 1)[0, 0]

    return sproj
//  [python]
def transfer_error(X, h):
    h_dims = h.size()
    batch_dims = list(h_dims)[:-1]
    H_dims = tuple(list(h_dims)[:-1] + [3, 3])

    N = X.size(-2)

    H = h.view(H_dims)
    I = torch.eye(3, device=X.device).view([1]*(len(H_dims)-2) + [3, 3])
    det = torch.det(H).view(batch_dims + [1, 1])
    H_ = torch.where(det < 1e-6, I, H)

    X1 = torch.ones(batch_dims + [N, 3, 1], device=X.device)
    X2 = torch.ones(batch_dims + [N, 3, 1], device=X.device)
    X1[..., 0:2, 0] = X[..., 0:2]
    X2[..., 0:2, 0] = X[..., 2:4]

    HX1 = H.unsqueeze(-3) @ X1
    HX1[..., 0, 0] /= torch.clamp(HX1[..., 2, 0], min=1e-8)
    HX1[..., 1, 0] /= torch.clamp(HX1[..., 2, 0], min=1e-8)
    HX1[..., 2, 0] /= torch.clamp(HX1[..., 2, 0], min=1e-8)

    HX2 = torch.solve(X2, H_.unsqueeze(-3))[0]
    HX2[..., 0, 0] /= torch.clamp(HX2[..., 2, 0], min=1e-8)
    HX2[..., 1, 0] /= torch.clamp(HX2[..., 2, 0], min=1e-8)
    HX2[..., 2, 0] /= torch.clamp(HX2[..., 2, 0], min=1e-8)

    signed_distances_1 = HX1 - X2
    distances_1 = (signed_distances_1 * signed_distances_1).sum(dim=-2).squeeze(-1)
    signed_distances_2 = HX2 - X1
    distances_2 = (signed_distances_2 * signed_distances_2).sum(dim=-2).squeeze(-1)

    distances = distances_1 + distances_2

    return distances
//  [python]
def transfer_error(X, h):
    h_dims = h.size()
    batch_dims = list(h_dims)[:-1]
    H_dims = tuple(list(h_dims)[:-1] + [3, 3])

    N = X.size(-2)

    H = h.view(H_dims)
    I = torch.eye(3, device=X.device).view([1]*(len(H_dims)-2) + [3, 3])
    det = torch.det(H).view(batch_dims + [1, 1])
    H_ = torch.where(det < 1e-6, I, H)

    X1 = torch.ones(batch_dims + [N, 3, 1], device=X.device)
    X2 = torch.ones(batch_dims + [N, 3, 1], device=X.device)
    X1[..., 0:2, 0] = X[..., 0:2]
    X2[..., 0:2, 0] = X[..., 2:4]

    HX1 = H.unsqueeze(-3) @ X1
    HX1[..., 0, 0] /= torch.clamp(HX1[..., 2, 0], min=1e-8)
    HX1[..., 1, 0] /= torch.clamp(HX1[..., 2, 0], min=1e-8)
    HX1[..., 2, 0] /= torch.clamp(HX1[..., 2, 0], min=1e-8)

    HX2 = torch.linalg.solve(H_.unsqueeze(-3), X2)
    HX2[..., 0, 0] /= torch.clamp(HX2[..., 2, 0], min=1e-8)
    HX2[..., 1, 0] /= torch.clamp(HX2[..., 2, 0], min=1e-8)
    HX2[..., 2, 0] /= torch.clamp(HX2[..., 2, 0], min=1e-8)

    signed_distances_1 = HX1 - X2
    distances_1 = (signed_distances_1 * signed_distances_1).sum(dim=-2).squeeze(-1)
    signed_distances_2 = HX2 - X1
    distances_2 = (signed_distances_2 * signed_distances_2).sum(dim=-2).squeeze(-1)

    distances = distances_1 + distances_2

    return distances
//  [python]
def ldet_sw_factor_analysis(Lambda, W, I_l=None):
    """
    Compute the log determinant of the covariance matrix in a factor
    analysis model using the Sherman Woodbury identity

    Parameters
    ----------
    Lambda : torch.Tensor
        Diagonal component of covariance matrix

    W : torch.Tensor
        Factor loading matrix.

    I_l : torch.Tensor, optional
        LxL identity matrix. If not provided, it is set to the identity matrix.

    Returns
    -------
    torch.Tensor
        Log determinant of the precision matrix.
    """
    if I_l is None:
        I_l = torch.eye(W.shape[0])
    ldet_L = torch.logdet(Lambda)
    LiW = torch.solve(torch.transpose(W, 0, 1), Lambda)[0]
    WtLiW = torch.matmul(W, LiW)
    IWtLiW = I_l + WtLiW
    ldet_p = torch.logdet(IWtLiW)
    log_determinant = ldet_p + ldet_L
    return log_determinant
//  [python]
def ldet_sw_factor_analysis(Lambda, W, I_l=None):
    """
    Compute the log determinant of the covariance matrix in a factor
    analysis model using the Sherman Woodbury identity

    Parameters
    ----------
    Lambda : torch.Tensor
        Diagonal component of covariance matrix

    W : torch.Tensor
        Factor loading matrix.

    I_l : torch.Tensor, optional
        LxL identity matrix. If not provided, it is set to the identity matrix.

    Returns
    -------
    torch.Tensor
        Log determinant of the precision matrix.
    """
    if I_l is None:
        I_l = torch.eye(W.shape[0])
    ldet_L = torch.logdet(Lambda)
    LiW = la.solve(Lambda, torch.transpose(W, 0, 1))
    WtLiW = torch.matmul(W, LiW)
    IWtLiW = I_l + WtLiW
    ldet_p = torch.logdet(IWtLiW)
    log_determinant = ldet_p + ldet_L
    return log_determinant
//  [python]
def inverse_sw_full(A, U, B, V):
    """
    Compute the inverse of a matrix product.

    M = (A + UBV)

    Parameters
    ----------
    A : torch.Tensor
        Input matrix A.

    U : torch.Tensor
        Input matrix U.

    B : torch.Tensor
        Input matrix B.

    V : torch.Tensor
        Input matrix V.

    Returns
    -------
    torch.Tensor
        Inverse of the matrix product.
    """
    # Compute the inverse of matrix A
    A_inv = la.inv(A)

    # Compute the inverse of matrix B
    B_inv = la.inv(B)

    # Compute the product A_inv * U
    AiU = torch.matmul(A_inv, U)

    # Compute the product V * A_inv
    VAinv = torch.matmul(V, A_inv)

    # Compute the product V * A_inv * U
    VAiU = torch.matmul(VAinv, U)

    # Compute the sum of B_inv and V * A_inv * U
    middle = B_inv + VAiU

    # Solve the system of linear equations using the product V * A_inv
    end = torch.solve(VAinv, middle)[0]

    # Compute the product A_inv * U * (solved system)
    second_term = torch.matmul(AiU, end)

    # Compute the inverse of the matrix product
    Sigma_inv = A_inv - second_term

    return Sigma_inv
//  [python]
def inverse_sw_full(A, U, B, V):
    """
    Compute the inverse of a matrix product.

    M = (A + UBV)

    Parameters
    ----------
    A : torch.Tensor
        Input matrix A.

    U : torch.Tensor
        Input matrix U.

    B : torch.Tensor
        Input matrix B.

    V : torch.Tensor
        Input matrix V.

    Returns
    -------
    torch.Tensor
        Inverse of the matrix product.
    """
    # Compute the inverse of matrix A
    A_inv = la.inv(A)

    # Compute the inverse of matrix B
    B_inv = la.inv(B)

    # Compute the product A_inv * U
    AiU = torch.matmul(A_inv, U)

    # Compute the product V * A_inv
    VAinv = torch.matmul(V, A_inv)

    # Compute the product V * A_inv * U
    VAiU = torch.matmul(VAinv, U)

    # Compute the sum of B_inv and V * A_inv * U
    middle = B_inv + VAiU

    # Solve the system of linear equations using the product V * A_inv
    end = la.solve(middle, VAinv)

    # Compute the product A_inv * U * (solved system)
    second_term = torch.matmul(AiU, end)

    # Compute the inverse of the matrix product
    Sigma_inv = A_inv - second_term

    return Sigma_inv
//  [python]
def _torch_solve_cast(A: Tensor, B: Tensor) -> Tensor:
    """Helper function to make torch.solve work with other than fp32/64.

    For stable operation, the input matrices should be cast to fp64, and the output will be cast back to the input
    dtype.
    """
    # cast to fp64 and solve
    out = torch.solve(A.to(torch.float64), B.to(torch.float64))[0]

    # cast back to the input dtype
    return out.to(A.dtype)
//  [python]
def _torch_solve_cast(A: Tensor, B: Tensor) -> Tensor:
    """Helper function to make torch.solve work with other than fp32/64.

    For stable operation, the input matrices should be cast to fp64, and the output will be cast back to the input
    dtype.
    """
    # cast to fp64 and solve
    out = torch.linalg.solve(A.to(torch.float64), B.to(torch.float64))

    # cast back to the input dtype
    return out.to(A.dtype)
//  [python]
def solve_output(self, Hidden_States, y, reg_lambda=1.0):
        '''
        sovle the output with ridge regression, reg_lambda (default) = 1.0
        '''
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self.device), Hidden_States), dim=1)


        col = t_hs.size(1)
        HTH = torch.mm(t_hs.t(), t_hs)
        HTY = torch.mm(t_hs.t(), y)

        I = (reg_lambda * torch.eye(HTH.size(0))).to(self.device)
        A = HTH + I

        orig_rank = torch.linalg.matrix_rank(A.float()).item()

        tag = 'Inverse' if orig_rank == col else 'Pseudo-inverse'

        assert tag == 'Inverse'
        
        if tag == 'Inverse':
            W = torch.solve(HTY, A)[0]
        else:
            W = torch.mm(torch.pinverse(A.cpu()),HTY.cpu()).to(self.device)
            
        # self.logger.info('Solving Method: {} \t L2 regular: {}'.format(tag, 'True' if self.reg_lambda != 0 else 'False'))
        return W, tag
//  [python]
def solve_output(self, Hidden_States, y, reg_lambda=1.0):
        '''
        sovle the output with ridge regression, reg_lambda (default) = 1.0
        '''
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self.device), Hidden_States), dim=1)


        col = t_hs.size(1)
        HTH = torch.mm(t_hs.t(), t_hs)
        HTY = torch.mm(t_hs.t(), y)

        I = (reg_lambda * torch.eye(HTH.size(0))).to(self.device)
        A = HTH + I

        orig_rank = torch.linalg.matrix_rank(A.float()).item()

        tag = 'Inverse' if orig_rank == col else 'Pseudo-inverse'

        assert tag == 'Inverse'
        
        if tag == 'Inverse':
            W = torch.linalg.solve(A, HTY)
        else:
            W = torch.mm(torch.linalg.pinv(A.cpu()),HTY.cpu()).to(self.device)
            
        # self.logger.info('Solving Method: {} \t L2 regular: {}'.format(tag, 'True' if self.reg_lambda != 0 else 'False'))
        return W, tag
//  [python]
def newton_raphson(self, numCoalEv, wNative, gamma, precision_matrix):
        iteration = 0
        grad = torch.tensor(float("inf"))
        gamma = gamma.clone()
        while (
            torch.linalg.vector_norm(grad) > self._stop_value
            and iteration < self._max_iterations
        ):
            jac = self.jacobian(wNative, gamma, precision_matrix)
            grad = self.gradient(numCoalEv, wNative, gamma, precision_matrix)
            gamma = gamma + torch.solve(grad.unsqueeze(1), jac)[0].squeeze(1)
            iteration += 1
        return gamma
//  [python]
def newton_raphson(self, numCoalEv, wNative, gamma, precision_matrix):
        iteration = 0
        grad = torch.tensor(float("inf"))
        gamma = gamma.clone()
        while (
            torch.linalg.vector_norm(grad) > self._stop_value
            and iteration < self._max_iterations
        ):
            jac = self.jacobian(wNative, gamma, precision_matrix)
            grad = self.gradient(numCoalEv, wNative, gamma, precision_matrix)
            gamma = gamma + torch.linalg.solve(jac, grad)
            iteration += 1
        return gamma
//  [python]
def mvn_log_prob_sw(X, mu, Lambda, W, I_l=None):
    """
    Compute the log probability of data points under a multivariate
    normal distribution using the Sherman Woodbury matrix identity.

    Parameters
    ----------
    X : torch.Tensor
        Data points (rows are observations, columns are variables).

    mu : torch.Tensor
        Mean vector of the multivariate normal distribution.

    Lambda : torch.Tensor
        Diagonal component of covariance matrix

    W : torch.Tensor
        Weight matrix.

    I_l : torch.Tensor, optional
        LxL identity matrix. If not provided, it is set to the
        identity matrix.

    Returns
    -------
    torch.Tensor
        Log probability of the data points under the multivariate
        normal distribution.
    """
    L, p = W.shape

    if I_l is None:
        I_l = torch.eye(L)

    term1 = -p * 0.9189385332046727

    X_demeaned = X - mu

    Wt = torch.transpose(W, 0, 1)

    F = torch.solve(Wt, Lambda)[0]
    E2 = I_l + W @ F

    ldet_L = torch.logdet(Lambda)
    ldet_E2 = torch.logdet(E2)
    term2 = -0.5 * (ldet_L + ldet_E2)

    end_center = torch.solve(F.T, E2)[0]
    middle = torch.inverse(Lambda) - F @ end_center
    end = middle @ X_demeaned.T

    quad = -1 * X_demeaned.T * end
    term3 = torch.sum(quad, axis=0)  # type: ignore

    log_prob_window = term1 + term2 + term3 / 2
    log_prob = torch.mean(log_prob_window)

    return log_prob
//  [python]
def mvn_log_prob_sw(X, mu, Lambda, W, I_l=None):
    """
    Compute the log probability of data points under a multivariate
    normal distribution using the Sherman Woodbury matrix identity.

    Parameters
    ----------
    X : torch.Tensor
        Data points (rows are observations, columns are variables).

    mu : torch.Tensor
        Mean vector of the multivariate normal distribution.

    Lambda : torch.Tensor
        Diagonal component of covariance matrix

    W : torch.Tensor
        Weight matrix.

    I_l : torch.Tensor, optional
        LxL identity matrix. If not provided, it is set to the
        identity matrix.

    Returns
    -------
    torch.Tensor
        Log probability of the data points under the multivariate
        normal distribution.
    """
    L, p = W.shape

    if I_l is None:
        I_l = torch.eye(L)

    term1 = -p * 0.9189385332046727

    X_demeaned = X - mu

    Wt = torch.transpose(W, 0, 1)

    F = la.solve(Lambda, Wt)
    E2 = I_l + W @ F

    ldet_L = torch.logdet(Lambda)
    ldet_E2 = torch.logdet(E2)
    term2 = -0.5 * (ldet_L + ldet_E2)

    end_center = la.solve(E2, F.T)
    middle = la.inv(Lambda) - F @ end_center
    end = middle @ X_demeaned.T

    quad = -1 * X_demeaned.T * end
    term3 = torch.sum(quad, axis=0)  # type: ignore

    log_prob_window = term1 + term2 + term3 / 2
    log_prob = torch.mean(log_prob_window)

    return log_prob
//  [python]
def optimise_ridge_parameter(self, Features, Y):
        ridges = 10.0 ** np.arange(-8, 9)
        num_val_samples = int(Features.shape[0] * 0.8)
        losses = []
        Q_val = Features[0:num_val_samples, :].T @ Y[0:num_val_samples, :]
        G_val = Features[0:num_val_samples, :].T @ Features[0:num_val_samples, :]
        for ridge in ridges:
            Wo = torch.solve(Q_val, G_val + ridge*torch.eye(G_val.size(dim=0)))[0].T #better nmerical stability than .inv
            Y_train_pred = Features[num_val_samples::,:] @ Wo.T
            losses.append(F.mse_loss(Y_train_pred, Y[num_val_samples::, :]))
        ridge = ridges[np.argmin(np.array(losses))]
        print('selected lambda =',ridge)
        return ridge
//  [python]
def optimise_ridge_parameter(self, Features, Y):
        ridges = 10.0 ** np.arange(-8, 9)
        num_val_samples = int(Features.shape[0] * 0.8)
        losses = []
        Q_val = Features[0:num_val_samples, :].T @ Y[0:num_val_samples, :]
        G_val = Features[0:num_val_samples, :].T @ Features[0:num_val_samples, :]
        for ridge in ridges:
            Wo = torch.linalg.solve(G_val + ridge*torch.eye(G_val.size(dim=0)), Q_val).T #better nmerical stability than .inv
            Y_train_pred = Features[num_val_samples::,:] @ Wo.T
            losses.append(F.mse_loss(Y_train_pred, Y[num_val_samples::, :]))
        ridge = ridges[np.argmin(np.array(losses))]
        print('selected lambda =',ridge)
        return ridge
//  [python]
def replace_fc(self, trainloader, model, args):       
        model = model.eval()
        embedding_list = []
        label_list = []
        with torch.no_grad():
            for i, batch in enumerate(trainloader):
                (_,data, label) = batch
                data = data.to(self._device)
                label = label.to(self._device)
                embedding = model.extract_vector(data)
                embedding_list.append(embedding.cpu())
                label_list.append(label.cpu())
        embedding_list = torch.cat(embedding_list, dim=0)
        label_list = torch.cat(label_list, dim=0)
        
        Y = target2onehot(label_list, self.args["nb_classes"])
        Features_h = F.relu(embedding_list @ self.W_rand.cpu())
        self.Q = self.Q + Features_h.T @ Y
        self.G = self.G + Features_h.T @ Features_h
        ridge = self.optimise_ridge_parameter(Features_h, Y)
        Wo = torch.inverse(self.G + ridge*torch.eye(self.G.size(dim=0))) @ self.Q.T # better nmerical stability than .invv
        self._network.fc.weight.data = Wo[0:self._network.fc.weight.shape[0],:].to(self._device)
        
        return model
//  [python]
def replace_fc(self, trainloader, model, args):       
        model = model.eval()
        embedding_list = []
        label_list = []
        with torch.no_grad():
            for i, batch in enumerate(trainloader):
                (_,data, label) = batch
                data = data.to(self._device)
                label = label.to(self._device)
                embedding = model.extract_vector(data)
                embedding_list.append(embedding.cpu())
                label_list.append(label.cpu())
        embedding_list = torch.cat(embedding_list, dim=0)
        label_list = torch.cat(label_list, dim=0)
        
        Y = target2onehot(label_list, self.args["nb_classes"])
        Features_h = F.relu(embedding_list @ self.W_rand.cpu())
        self.Q = self.Q + Features_h.T @ Y
        self.G = self.G + Features_h.T @ Features_h
        ridge = self.optimise_ridge_parameter(Features_h, Y)
        Wo = torch.linalg.solve(self.G + ridge*torch.eye(self.G.size(dim=0)), self.Q).T # better nmerical stability than .invv
        self._network.fc.weight.data = Wo[0:self._network.fc.weight.shape[0],:].to(self._device)
        
        return model
//  [python]
def solve(mat: torch.tensor, vec: torch.tensor,
              *args, **kwargs) -> torch.tensor:
        return torch.solve(vec, mat)[0]
//  [python]
def solve(mat: torch.tensor, vec: torch.tensor,
              *args, **kwargs) -> torch.tensor:
        return torch.linalg.solve(mat, vec)
//  [python]
def solve(y, d, smoothing, kernel, epsilon, powers):
    """Build then solve the RBF linear system"""

    lhs, rhs, shift, scale = build(y, d, smoothing, kernel, epsilon, powers)
    try:
        coeffs = torch.solve(rhs, lhs)[0]
    except RuntimeError:  # singular matrix
        if coeffs is None:
            msg = "Singular matrix."
            nmonos = powers.shape[0]
            if nmonos > 0:
                pmat = polynomial_matrix((y - shift) / scale, powers)
                rank = torch.linalg.matrix_rank(pmat)
                if rank < nmonos:
                    msg = (
                        "Singular matrix. The matrix of monomials evaluated at "
                        "the data point coordinates does not have full column "
                        f"rank ({rank}/{nmonos})."
                    )

            raise ValueError(msg)

    return shift, scale, coeffs
//  [python]
def solve(y, d, smoothing, kernel, epsilon, powers):
    """Build then solve the RBF linear system"""

    lhs, rhs, shift, scale = build(y, d, smoothing, kernel, epsilon, powers)
    try:
        coeffs = torch.linalg.solve(lhs, rhs)
    except RuntimeError:  # singular matrix
        if coeffs is None:
            msg = "Singular matrix."
            nmonos = powers.shape[0]
            if nmonos > 0:
                pmat = polynomial_matrix((y - shift) / scale, powers)
                rank = torch.linalg.matrix_rank(pmat)
                if rank < nmonos:
                    msg = (
                        "Singular matrix. The matrix of monomials evaluated at "
                        "the data point coordinates does not have full column "
                        f"rank ({rank}/{nmonos})."
                    )

            raise ValueError(msg)

    return shift, scale, coeffs
//  [python]
def forward(self, covmat, rets=None):
        """Perform forward pass.

        Parameters
        ----------
        covmat : torch.Tensor
            Covariance matrix of shape `(n_samples, n_assets, n_assets)`.

        rets : torch.Tensor or None
            If tensor then of shape `(n_samples, n_assets)` representing expected returns. If provided triggers
            computation of maximum share ratio. Else None triggers computation of minimum variance portfolio.

        Returns
        -------
        weights : torch.Tensor
            Of shape (n_samples, n_assets) representing the optimal weights. If `rets` provided, then it represents
            maximum sharpe ratio portfolio (tangency portfolio). Otherwise minimum variance portfolio.
        """
        n_samples, n_assets, _ = covmat.shape
        device = covmat.device
        dtype = covmat.dtype

        ones = torch.ones(n_samples, n_assets, 1).to(
            device=device, dtype=dtype
        )
        if rets is not None:
            expected_returns = rets.view(n_samples, n_assets, 1)
        else:
            expected_returns = ones

        w_unscaled = torch.solve(expected_returns, covmat)[0]
        denominator = torch.matmul(ones.permute(0, 2, 1), w_unscaled)
        w = w_unscaled / denominator

        return w.squeeze(-1)
//  [python]
def forward(self, covmat, rets=None):
        """Perform forward pass.

        Parameters
        ----------
        covmat : torch.Tensor
            Covariance matrix of shape `(n_samples, n_assets, n_assets)`.

        rets : torch.Tensor or None
            If tensor then of shape `(n_samples, n_assets)` representing expected returns. If provided triggers
            computation of maximum share ratio. Else None triggers computation of minimum variance portfolio.

        Returns
        -------
        weights : torch.Tensor
            Of shape (n_samples, n_assets) representing the optimal weights. If `rets` provided, then it represents
            maximum sharpe ratio portfolio (tangency portfolio). Otherwise minimum variance portfolio.
        """
        n_samples, n_assets, _ = covmat.shape
        device = covmat.device
        dtype = covmat.dtype

        ones = torch.ones(n_samples, n_assets, 1).to(
            device=device, dtype=dtype
        )
        if rets is not None:
            expected_returns = rets.view(n_samples, n_assets, 1)
        else:
            expected_returns = ones

        w_unscaled = torch.linalg.solve(covmat, expected_returns)
        denominator = torch.matmul(ones.permute(0, 2, 1), w_unscaled)
        w = w_unscaled / denominator

        return w.squeeze(-1)
//  [python]
def schur_solve(cam_block: torch.Tensor, cross_block: torch.Tensor, point_block: torch.Tensor, g_cam: torch.Tensor, g_points: torch.Tensor):
    N, D, _ = point_block.shape # N, D, D
    M, th, _ = cam_block.shape # M, th, th
    K, _, _, _ = cross_block.shape # K, M, th, D
    g_cam.shape # M, th
    
    
    assert N == K, "Assuming K is equal to N"
    point_block_inv = torch.linalg.inv(point_block) # N, D, D
    rhs = einsum(point_block_inv, cross_block, "k d1 d2, k m2 th2 d2 -> k m2 th2 d1")
    BTA_invB = einsum(cross_block, rhs, "k m1 th1 d1, k m2 th2 d1 -> m1 th1 m2 th2").reshape(M*th, M*th)
    S = torch.block_diag(*cam_block) - BTA_invB
    
    part_of_camera_rhs = einsum(point_block_inv, g_points,"n d1 d2, n d2 -> n d1")
    camera_rhs = (g_cam - einsum(cross_block, part_of_camera_rhs, "n m th d1, n d1 -> m th")).reshape(M*th)
    camera_params = torch.solve(camera_rhs, S)[0].reshape(M, th)
    
    points_rhs = g_points - einsum(cross_block, camera_params, "n m th d2, m th -> n d2")
    point_coordinates = einsum(point_block_inv, points_rhs,"n d1 d2, n d2 -> n d1")
    return camera_params, point_coordinates
//  [python]
def schur_solve(cam_block: torch.Tensor, cross_block: torch.Tensor, point_block: torch.Tensor, g_cam: torch.Tensor, g_points: torch.Tensor):
    N, D, _ = point_block.shape # N, D, D
    M, th, _ = cam_block.shape # M, th, th
    K, _, _, _ = cross_block.shape # K, M, th, D
    g_cam.shape # M, th
    
    
    assert N == K, "Assuming K is equal to N"
    point_block_inv = torch.linalg.inv(point_block) # N, D, D
    rhs = einsum(point_block_inv, cross_block, "k d1 d2, k m2 th2 d2 -> k m2 th2 d1")
    BTA_invB = einsum(cross_block, rhs, "k m1 th1 d1, k m2 th2 d1 -> m1 th1 m2 th2").reshape(M*th, M*th)
    S = torch.block_diag(*cam_block) - BTA_invB
    
    part_of_camera_rhs = einsum(point_block_inv, g_points,"n d1 d2, n d2 -> n d1")
    camera_rhs = (g_cam - einsum(cross_block, part_of_camera_rhs, "n m th d1, n d1 -> m th")).reshape(M*th)
    camera_params = torch.linalg.solve(S, camera_rhs).reshape(M, th)
    
    points_rhs = g_points - einsum(cross_block, camera_params, "n m th d2, m th -> n d2")
    point_coordinates = einsum(point_block_inv, points_rhs,"n d1 d2, n d2 -> n d1")
    return camera_params, point_coordinates
//  [python]
def fit(self, X: Tensor, y: Tensor):
        assert len(X.shape) == 3
        assert len(y.shape) == 2
        self.X_ = X
        self.y_ = y

        k = self.kernel(X, X)
        # klmbda = k + torch.eye(k.shape[0]) * self.lmbda
        # Below is the same as above but avoids the creation of a new tensor on a different device
        klmbda = (k - k.diag().diag()) + (self.lmbda * self.X_.shape[0] + k.diag()).diag()
        self.k_ = k
        self.alpha_ = torch.solve(y, klmbda)[0]
        return self
//  [python]
def fit(self, X: Tensor, y: Tensor):
        assert len(X.shape) == 3
        assert len(y.shape) == 2
        self.X_ = X
        self.y_ = y

        k = self.kernel(X, X)
        # klmbda = k + torch.eye(k.shape[0]) * self.lmbda
        # Below is the same as above but avoids the creation of a new tensor on a different device
        klmbda = (k - k.diag().diag()) + (self.lmbda * self.X_.shape[0] + k.diag()).diag()
        self.k_ = k
        self.alpha_ = LA.solve(klmbda, y)
        return self
//  [python]
def update_axiliary_variable(self):
        self.Yinv_FTMM = torch.solve(
            torch.einsum("nft, nfij -> ftij", self.PSD_NFT.to(self.TYPE_COMPLEX), self.SCM_NFMM),
            torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device)[None, None],
        )[1]

        Yinv_x_FTM = torch.einsum("ftij, ftj -> fti", self.Yinv_FTMM, self.X_FTM)
        self.Yinv_X_Yinv_FTMM = torch.einsum("fti, ftj -> ftij", Yinv_x_FTM, Yinv_x_FTM.conj())
        self.tr_SCM_Yinv_X_Yinv_NFT = torch_trace(
            torch.einsum("nfij, ftjl -> nftil", self.SCM_NFMM, self.Yinv_X_Yinv_FTMM)
        ).real
        self.tr_SCM_Yinv_NFT = torch_trace(torch.einsum("nfij, ftjl -> nftil", self.SCM_NFMM, self.Yinv_FTMM)).real
//  [python]
def update_axiliary_variable(self):
        self.Yinv_FTMM = torch.linalg.solve(
            torch.einsum("nft, nfij -> ftij", self.PSD_NFT.to(self.TYPE_COMPLEX), self.SCM_NFMM),
            torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device)[None, None],
        )

        Yinv_x_FTM = torch.einsum("ftij, ftj -> fti", self.Yinv_FTMM, self.X_FTM)
        self.Yinv_X_Yinv_FTMM = torch.einsum("fti, ftj -> ftij", Yinv_x_FTM, Yinv_x_FTM.conj())
        self.tr_SCM_Yinv_X_Yinv_NFT = torch_trace(
            torch.einsum("nfij, ftjl -> nftil", self.SCM_NFMM, self.Yinv_X_Yinv_FTMM)
        ).real
        self.tr_SCM_Yinv_NFT = torch_trace(torch.einsum("nfij, ftjl -> nftil", self.SCM_NFMM, self.Yinv_FTMM)).real
//  [python]
if operator == 'symmetric':
    Ls = symmetric(idx, val)
    Pr = torch.matrix_power(torch.eye(kernel.base_kernel.nodes.shape[0])*2*nu/lengthscale.square().squeeze() + Ls, nu)
elif operator == 'randomwalk':
    Lr, D = randomwalk(idx, val)
    Pr = torch.matrix_power(torch.eye(kernel.base_kernel.nodes.shape[0])*2*nu/lengthscale.square().squeeze() + Lr, nu)
    Pr = torch.mm(D.diag(), Pr)

# Semisupervised
if semisupervised:
    Qxx = Pr[:, labeled]
    Qxx = Qxx[labeled, :]
    Qxz = Pr[:, not_labeled]
    Qxz = Qxz[labeled, :]
    Qzz = Pr[:, not_labeled]
    Qzz = Qzz[not_labeled, :]
    Qzx = Pr[:, labeled]
    Qzx = Qzx[not_labeled, :]
    Pr = Qxx - torch.mm(Qxz, torch.solve(Qzx, Qzz)[0])
//  [python]
if operator == 'symmetric':
    Ls = symmetric(idx, val)
    Pr = torch.matrix_power(torch.eye(kernel.base_kernel.nodes.shape[0])*2*nu/lengthscale.square().squeeze() + Ls, nu)
elif operator == 'randomwalk':
    Lr, D = randomwalk(idx, val)
    Pr = torch.matrix_power(torch.eye(kernel.base_kernel.nodes.shape[0])*2*nu/lengthscale.square().squeeze() + Lr, nu)
    Pr = torch.mm(D.diag(), Pr)

# Semisupervised
if semisupervised:
    Qxx = Pr[:, labeled]
    Qxx = Qxx[labeled, :]
    Qxz = Pr[:, not_labeled]
    Qxz = Qxz[labeled, :]
    Qzz = Pr[:, not_labeled]
    Qzz = Qzz[not_labeled, :]
    Qzx = Pr[:, labeled]
    Qzx = Qzx[not_labeled, :]
    Pr = Qxx - torch.mm(Qxz, torch.linalg.solve(Qzz, Qzx))
//  [python]
def proju_original(x, u):
    import geoopt.linalg as linalg

    # takes batch data
    # batch_size, n, _ = x.shape
    x_shape = x.shape
    x = x.reshape(-1, x_shape[-2], x_shape[-1])
    batch_size, n = x.shape[0:2]

    e = torch.ones(batch_size, n, 1, dtype=x.dtype)
    I = torch.unsqueeze(torch.eye(x.shape[-1], dtype=x.dtype), 0).repeat(
        batch_size, 1, 1
    )

    mu = x * u

    A = linalg.block_matrix([[I, x], [torch.transpose(x, 1, 2), I]])

    B = A[:, :, 1:]
    b = torch.cat(
        [
            torch.sum(mu, dim=2, keepdim=True),
            torch.transpose(torch.sum(mu, dim=1, keepdim=True), 1, 2),
        ],
        dim=1,
    )

    zeta = torch.solve(
        B.transpose(1, 2) @ B, B.transpose(1, 2) @ (b - A[:, :, 0:1])
    )[0]
    alpha = torch.cat(
        [torch.ones(batch_size, 1, 1, dtype=x.dtype), zeta[:, 0 : n - 1]], dim=1
    )
    beta = zeta[:, n - 1 : 2 * n - 1]
    rgrad = mu - (alpha @ e.transpose(1, 2) + e @ beta.transpose(1, 2)) * x

    rgrad = rgrad.reshape(x_shape)
    return rgrad
//  [python]
def proju_original(x, u):
    import geoopt.linalg as linalg

    # takes batch data
    # batch_size, n, _ = x.shape
    x_shape = x.shape
    x = x.reshape(-1, x_shape[-2], x_shape[-1])
    batch_size, n = x.shape[0:2]

    e = torch.ones(batch_size, n, 1, dtype=x.dtype)
    I = torch.unsqueeze(torch.eye(x.shape[-1], dtype=x.dtype), 0).repeat(
        batch_size, 1, 1
    )

    mu = x * u

    A = linalg.block_matrix([[I, x], [torch.transpose(x, 1, 2), I]])

    B = A[:, :, 1:]
    b = torch.cat(
        [
            torch.sum(mu, dim=2, keepdim=True),
            torch.transpose(torch.sum(mu, dim=1, keepdim=True), 1, 2),
        ],
        dim=1,
    )

    zeta = torch.linalg.solve(
        B.transpose(1, 2) @ B, B.transpose(1, 2) @ (b - A[:, :, 0:1])
    )
    alpha = torch.cat(
        [torch.ones(batch_size, 1, 1, dtype=x.dtype), zeta[:, 0 : n - 1]], dim=1
    )
    beta = zeta[:, n - 1 : 2 * n - 1]
    rgrad = mu - (alpha @ e.transpose(1, 2) + e @ beta.transpose(1, 2)) * x

    rgrad = rgrad.reshape(x_shape)
    return rgrad
//  [python]
def _step(self) -> Tensor:
        coalescent = self.coalescent.distribution()
        gamma = self.gmrf.field.tensor
        sufficient_statistics, coalescent_counts = coalescent.sufficient_statistics(
            self.coalescent.tree_model.node_heights
        )
        precision_matrix = self.gmrf.precision_matrix()
        self.gmrf.precision.tensor = self.propose_precision()
        proposed_precision_matrix = self.gmrf.precision_matrix()

        mode_forward = self.newton_raphson(
            coalescent_counts, sufficient_statistics, gamma, proposed_precision_matrix
        )

        forwardQW = proposed_precision_matrix.clone()
        backwardQW = precision_matrix.clone()

        dim = precision_matrix.shape[-1]

        diagonal1 = sufficient_statistics * torch.exp(-mode_forward)
        forwardQW[range(dim), range(dim)] += diagonal1
        diagonal1 = diagonal1 * (mode_forward + 1) - coalescent_counts

        z = torch.randn(dim)

        try:
            cholesky = torch.linalg.cholesky(forwardQW, upper=True)
        except torch._C._LinAlgError:
            return torch.tensor(torch.inf)

        v = torch.solve(diagonal1.unsqueeze(1), cholesky.t())[0].squeeze(1)
        mu = torch.solve(v.unsqueeze(1), cholesky)[0].squeeze(1)
        u = torch.solve(z.unsqueeze(1), cholesky)[0].squeeze(1)
        proposed_gamma = mu + u

        self.gmrf.field.tensor = proposed_gamma.clone()

        diagonal = cholesky[range(dim), range(dim)]
        log_q_forward = diagonal[diagonal > 0.0000001].log().sum() - 0.5 * (z @ z)

        mode_backward = self.newton_raphson(
            coalescent_counts, sufficient_statistics, proposed_gamma, precision_matrix
        )

        diagonal1 = sufficient_statistics * torch.exp(-mode_backward)
        backwardQW[range(dim), range(dim)] += diagonal1
        diagonal1 = diagonal1 * (mode_backward + 1) - coalescent_counts

        try:
            cholesky = torch.linalg.cholesky(backwardQW, upper=True)
        except torch._C._LinAlgError:
            return torch.tensor(torch.inf)

        v = torch.solve(diagonal1.unsqueeze(1), cholesky.t())[0].squeeze(1)
        mu = torch.solve(v.unsqueeze(1), cholesky)[0].squeeze(1)
        diagonal1 = gamma - mu
        diagonal3 = backwardQW @ diagonal1

        diagonal = cholesky[range(dim), range(dim)]
        log_q_backward = diagonal[diagonal > 0.0000001].log().sum() - 0.5 * (
            diagonal1 @ diagonal3
        )
        return log_q_backward - log_q_forward
//  [python]
def _step(self) -> Tensor:
        coalescent = self.coalescent.distribution()
        gamma = self.gmrf.field.tensor
        sufficient_statistics, coalescent_counts = coalescent.sufficient_statistics(
            self.coalescent.tree_model.node_heights
        )
        precision_matrix = self.gmrf.precision_matrix()
        self.gmrf.precision.tensor = self.propose_precision()
        proposed_precision_matrix = self.gmrf.precision_matrix()

        mode_forward = self.newton_raphson(
            coalescent_counts, sufficient_statistics, gamma, proposed_precision_matrix
        )

        forwardQW = proposed_precision_matrix.clone()
        backwardQW = precision_matrix.clone()

        dim = precision_matrix.shape[-1]

        diagonal1 = sufficient_statistics * torch.exp(-mode_forward)
        forwardQW[range(dim), range(dim)] += diagonal1
        diagonal1 = diagonal1 * (mode_forward + 1) - coalescent_counts

        z = torch.randn(dim)

        try:
            cholesky = torch.linalg.cholesky(forwardQW, upper=True)
        except torch._C._LinAlgError:
            return torch.tensor(torch.inf)

        v = torch.linalg.solve(cholesky.t(), diagonal1)
        mu = torch.linalg.solve(cholesky, v)
        u = torch.linalg.solve(cholesky, z)
        proposed_gamma = mu + u

        self.gmrf.field.tensor = proposed_gamma.clone()

        diagonal = cholesky[range(dim), range(dim)]
        log_q_forward = diagonal[diagonal > 0.0000001].log().sum() - 0.5 * (z @ z)

        mode_backward = self.newton_raphson(
            coalescent_counts, sufficient_statistics, proposed_gamma, precision_matrix
        )

        diagonal1 = sufficient_statistics * torch.exp(-mode_backward)
        backwardQW[range(dim), range(dim)] += diagonal1
        diagonal1 = diagonal1 * (mode_backward + 1) - coalescent_counts

        try:
            cholesky = torch.linalg.cholesky(backwardQW, upper=True)
        except torch._C._LinAlgError:
            return torch.tensor(torch.inf)

        v = torch.linalg.solve(cholesky.t(), diagonal1)
        mu = torch.linalg.solve(cholesky, v)
        diagonal1 = gamma - mu
        diagonal3 = backwardQW @ diagonal1

        diagonal = cholesky[range(dim), range(dim)]
        log_q_backward = diagonal[diagonal > 0.0000001].log().sum() - 0.5 * (
            diagonal1 @ diagonal3
        )
        return log_q_backward - log_q_forward