code: def _conjugate_gradient(f_Ax, b, cg_iters, residual_tol=1e-10):
      """Use Conjugate Gradient iteration to solve Ax = b. Demmel p 312.
  
      Args:
          f_Ax (callable): A function to compute Hessian vector product.
          b (torch.Tensor): Right hand side of the equation to solve.
          cg_iters (int): Number of iterations to run conjugate gradient
              algorithm.
          residual_tol (float): Tolerence for convergence.
  
      Returns:
          torch.Tensor: Solution x* for equation Ax = b.
  
      """
      p = b.clone()
      r = b.clone()
      x = torch.zeros_like(b)
      rdotr = torch.dot(r, r)
  
      for _ in range(cg_iters):
          z = f_Ax(p)
          v = rdotr / torch.dot(p, z)
          x += v * p
          r -= v * z
          newrdotr = torch.dot(r, r)
          mu = newrdotr / rdotr
          p = r + mu * p
  
          rdotr = newrdotr
          if rdotr < residual_tol:
              break
      return x
  
file_path: garage/src/garage/torch/optimizers/conjugate_gradient_optimizer.py
imports: import warnings
  from dowel import logger
  import numpy
  import torch
  from torch.optim import Optimizer
  from garage.np import unflatten_tensors
label_desc: Use image node_id 2 for calling a global function with example usage: _conjugate_gradient(f_Ax, b, cg_iters, residual_tol) and returns: x