code: def cglb_conjugate_gradient(
      K: TensorType,
      b: TensorType,
      initial: TensorType,
      preconditioner: NystromPreconditioner,
      cg_tolerance: float,
      max_steps: int,
      restart_cg_step: int,
  ) -> tf.Tensor:
      """
      Conjugate gradient algorithm used in CGLB model. The method of
      conjugate gradient (Hestenes and Stiefel, 1952) produces a
      sequence of vectors :math:`v_0, v_1, v_2, ..., v_N` such that
      :math:`v_0` = initial, and (in exact arithmetic)
      :math:`Kv_n = b`. In practice, the v_i often converge quickly to
      approximate :math:`K^{-1}b`, and the algorithm can be stopped
      without running N iterations.
  
      We assume the preconditioner, :math:`Q`, satisfies :math:`Q ≺ K`,
      and stop the algorithm when :math:`r_i = b - Kv_i` satisfies
      :math:`||rᵢᵀ||_{Q⁻¹r}^2 = rᵢᵀQ⁻¹rᵢ <= ϵ`.
  
      :param K: Matrix we want to backsolve from. Must be PSD.
      :param b: Vector we want to backsolve.
      :param initial: Initial vector solution.
      :param preconditioner: Preconditioner function.
      :param cg_tolerance: Expected maximum error. This value is used
          as a decision boundary against stopping criteria.
      :param max_steps: Maximum number of CG iterations.
      :param restart_cg_step: Restart step at which the CG resets the
          internal state to the initial position using the currect
          solution vector :math:`v`. Can help avoid build up of
          numerical errors.
  
      :return: `v` where `v` approximately satisfies :math:`Kv = b`.
      """
  
      class CGState(NamedTuple):
          i: tf.Tensor
          v: tf.Tensor
          r: tf.Tensor
          p: tf.Tensor
          rz: tf.Tensor
  
      def stopping_criterion(state: CGState) -> tf.Tensor:
          return (0.5 * state.rz > cg_tolerance) and (
              state.i < max_steps
          )
  
      def cg_step(state: CGState) -> List[CGState]:
          Ap = state.p @ K
          denom = tf.reduce_sum(state.p * Ap, axis=-1)
          gamma = state.rz / denom
          v = state.v + gamma * state.p
          i = state.i + 1
          r = tf.cond(
              state.i % restart_cg_step == restart_cg_step - 1,
              lambda: b - v @ K,
              lambda: state.r - gamma * Ap,
          )
          z, new_rz = preconditioner(r)
          p = tf.cond(
              state.i % restart_cg_step == restart_cg_step - 1,
              lambda: z,
              lambda: z + state.p * new_rz / state.rz,
          )
          return [CGState(i, v, r, p, new_rz)]
  
      Kv = initial @ K
      r = b - Kv
      z, rz = preconditioner(r)
      p = z
      i = tf.constant(0, dtype=default_int())
      initial_state = CGState(i, initial, r, p, rz)
      final_state = tf.while_loop(
          stopping_criterion, cg_step, [initial_state]
      )
      final_state = tf.nest.map_structure(tf.stop_gradient, final_state)
      return final_state[0].v
  
imports: from typing import Any, List, NamedTuple, Optional, Tuple
  import tensorflow
  from check_shapes import check_shapes, inherit_check_shapes
  from ..base import InputData, MeanAndVariance, Parameter, RegressionData, TensorType
  from ..config import default_float, default_int
  from ..covariances import Kuf
  from ..utilities import add_noise_cov, assert_params_false, to_default_float
  from .sgpr import SGPR
label_desc: Use image node_id 1 for calling a global function with example usage: cglb_conjugate_gradient(K, b, initial, preconditioner, cg_tolerance, max_steps, restart_cg_step) and returns: final_state, unknown, list