class_docstr: Conjugate gradient algorithm which iteratively finds a solution $x$ for a system of linear  
  equations of the form $A x = b$, where $A x$ could be, for instance, a locally linear  
  approximation of a high-dimensional function.
  
  See below pseudo-code taken from  
  [Wikipedia](https://en.wikipedia.org/wiki/Conjugate_gradient_method#The_resulting_algorithm):
  
  ```text
  def conjgrad(A, b, x_0):
      r_0 := b - A * x_0
      c_0 := r_0
      r_0^2 := r^T * r
  
      for t in 0, ..., max_iterations - 1:
          Ac := A * c_t
          cAc := c_t^T * Ac
          lpha := r_t^2 / cAc
          x_{t+1} := x_t + lpha * c_t
          r_{t+1} := r_t - lpha * Ac
          r_{t+1}^2 := r_{t+1}^T * r_{t+1}
          if r_{t+1} < \epsilon:
              break
          eta = r_{t+1}^2 / r_t^2
          c_{t+1} := r_{t+1} + eta * c_t
  
      return x_{t+1}
  ```
class_docstr_tok: ["Conjugate","gradient","algorithm","which","iteratively","finds","a","solution","$","x","$","for","a","system","of","linear","equations","of","the","form","$","A","x","=","b","$",",","where","$","A","x","$","could","be",",","for","instance",",","a","locally","linear","approximation","of","a","high-dimensional","function",".","See","below","pseudo-code","taken","from","[","Wikipedia","]","(","https",":","\/\/en.wikipedia.org\/wiki\/Conjugate_gradient_method","#","The_resulting_algorithm",")",":","``","`","text","def","conjgrad","(","A",",","b",",","x_0",")",":","r_0",":","=","b","-","A","*","x_0","c_0",":","=","r_0","r_0^2",":","=","r^T","*","r","for","t","in","0",",","...",",","max_iterations","-","1",":","Ac",":","=","A","*","c_t","cAc",":","=","c_t^T","*","Ac","\u0007lpha",":","=","r_t^2","\/","cAc","x_","{","t+1","}",":","=","x_t","+","\u0007lpha","*","c_t","r_","{","t+1","}",":","=","r_t","-","\u0007lpha","*","Ac","r_","{","t+1","}","^2",":","=","r_","{","t+1","}","^T","*","r_","{","t+1","}","if","r_","{","t+1","}","<","\\epsilon",":","break","\beta","=","r_","{","t+1","}","^2","\/","r_t^2","c_","{","t+1","}",":","=","r_","{","t+1","}","+","\beta","*","c_t","return","x_","{","t+1","}","``","`"]
docstr: Iteration loop body of the conjugate gradient algorithm.
  
  Args:
      arguments: ???
      x: Current solution estimate $x_t$.
      conjugate: Current conjugate $c_t$.
      residual: Current residual $r_t$.
      squared_residual: Current squared residual $r_t^2$.
  
  Returns:
      Updated arguments for next iteration.
docstr_tok: ["Iteration","loop","body","of","the","conjugate","gradient","algorithm",".","Args",":","arguments",":","?","?","?","x",":","Current","solution","estimate","$","x_t","$",".","conjugate",":","Current","conjugate","$","c_t","$",".","residual",":","Current","residual","$","r_t","$",".","squared_residual",":","Current","squared","residual","$","r_t^2","$",".","Returns",":","Updated","arguments","for","next","iteration","."]
code: def step(
      self, *, arguments, x, conjugate, residual, squared_residual
  ):
      """
      Iteration loop body of the conjugate gradient algorithm.
  
      Args:
          arguments: ???
          x: Current solution estimate $x_t$.
          conjugate: Current conjugate $c_t$.
          residual: Current residual $r_t$.
          squared_residual: Current squared residual $r_t^2$.
  
      Returns:
          Updated arguments for next iteration.
      """
  
      # Ac := A * c_t
      A_conjugate = self.fn_x(arguments, conjugate)
  
      # TODO: reference?
      damping = self.damping.value()
  
      def no_damping():
          return A_conjugate
  
      def apply_damping():
          return A_conjugate.fmap(
              function=(lambda A_conj, conj: A_conj + damping * conj),
              zip_values=conjugate,
          )
  
      zero = tf_util.constant(value=0.0, dtype="float")
      skip_damping = tf.math.equal(x=damping, y=zero)
      A_conjugate = tf.cond(
          pred=skip_damping, true_fn=no_damping, false_fn=apply_damping
      )
  
      # cAc := c_t^T * Ac
      multiply = tf.math.multiply
      conjugate_A_conjugate = conjugate.fmap(
          function=multiply, zip_values=A_conjugate
      )
      conjugate_A_conjugate = tf.math.add_n(
          inputs=[
              tf.math.reduce_sum(input_tensor=conj_A_conj)
              for conj_A_conj in conjugate_A_conjugate.values()
          ]
      )
  
      # \alpha := r_t^2 / cAc
      epsilon = tf_util.constant(value=util.epsilon, dtype="float")
      conjugate_A_conjugate = tf.where(
          condition=(conjugate_A_conjugate > 0.0),
          x=(conjugate_A_conjugate + epsilon),
          y=(conjugate_A_conjugate - epsilon),
      )
      alpha = squared_residual / conjugate_A_conjugate
  
      # x_{t+1} := x_t + \alpha * c_t
      next_x = x.fmap(
          function=(lambda t, conj: t + alpha * conj),
          zip_values=conjugate,
      )
  
      # r_{t+1} := r_t - \alpha * Ac
      next_residual = residual.fmap(
          function=(lambda res, A_conj: res - alpha * A_conj),
          zip_values=A_conjugate,
      )
  
      # r_{t+1}^2 := r_{t+1}^T * r_{t+1}
      next_squared_residual = tf.math.add_n(
          inputs=[
              tf.math.reduce_sum(input_tensor=(res * res))
              for res in next_residual.values()
          ]
      )
  
      # \beta = r_{t+1}^2 / r_t^2
      beta = next_squared_residual / (squared_residual + epsilon)
  
      # c_{t+1} := r_{t+1} + \beta * c_t
      next_conjugate = next_residual.fmap(
          function=(lambda res, conj: res + beta * conj),
          zip_values=conjugate,
      )
  
      return (
          arguments,
          next_x,
          next_conjugate,
          next_residual,
          next_squared_residual,
      )
  
file_path: tensorforce/tensorforce/core/optimizers/solvers/conjugate_gradient.py
imports: import functools
  import tensorflow
  from tensorforce import util
  from tensorforce.core import parameter_modules, SignatureDict, TensorSpec, tf_function, tf_util
  from tensorforce.core.optimizers.solvers import Iterative
label_desc: Use image node_id 6 for calling the ConjugateGradient obj's underlying member method code with example usage: obj.step() and returns: arguments,  next_x,  next_conjugate,  next_residual,  next_squared_residual, A_conjugate, A_conjugate