code: def _cg_solve(
      A, b, x0=None, *, maxiter, tol=1e-5, atol=0.0, M=_identity
  ):
      # tolerance handling uses the "non-legacy" behavior of scipy.sparse.linalg.cg
      bs = _vdot_real_tree(b, b)
      atol2 = jnp.maximum(jnp.square(tol) * bs, jnp.square(atol))
  
      # https://en.wikipedia.org/wiki/Conjugate_gradient_method#The_preconditioned_conjugate_gradient_method
  
      def cond_fun(value):
          _, r, gamma, _, k = value
          rs = gamma.real if M is _identity else _vdot_real_tree(r, r)
          return (rs > atol2) & (k < maxiter)
  
      def body_fun(value):
          x, r, gamma, p, k = value
          Ap = A(p)
          alpha = gamma / _vdot_real_tree(p, Ap).astype(dtype)
          x_ = _add(x, _mul(alpha, p))
          r_ = _sub(r, _mul(alpha, Ap))
          z_ = M(r_)
          gamma_ = _vdot_real_tree(r_, z_).astype(dtype)
          beta_ = gamma_ / gamma
          p_ = _add(z_, _mul(beta_, p))
          return x_, r_, gamma_, p_, k + 1
  
      r0 = _sub(b, A(x0))
      p0 = z0 = M(r0)
      dtype = jnp.result_type(*tree_leaves(p0))
      gamma0 = _vdot_real_tree(r0, z0).astype(dtype)
      initial_value = (x0, r0, gamma0, p0, 0)
  
      x_final, *_ = lax.while_loop(cond_fun, body_fun, initial_value)
  
      return x_final
  
imports: from functools import partial
  import operator
  import numpy
  import jax
  import jax.numpy
  from jax import device_put
  from jax import lax
  from jax import scipy
  from jax.tree_util import tree_leaves, tree_map, tree_structure, tree_reduce, Partial
  from jax._src import dtypes
  from jax._src.lax import lax
  from jax._src.util import safe_map
label_desc: Use image node_id 9 for calling a global function with example usage: _cg_solve(A, b, x0) and returns: x_final, unknown, x_,  r_,  gamma_,  p_,  unknown