code: def conjugate_gradient(
      a_mat: types.Tensor, b_mat: types.Tensor, tol: float = 1e-10
  ) -> types.Float:
      """Returns `X` such that `A * X = B`.
  
      Implements the Conjugate Gradient method.
      https://en.wikipedia.org/wiki/Conjugate_gradient_method
  
      Args:
        a_mat: a Symmetric Positive Definite matrix, represented as a `Tensor` of
          shape `[n, n]`.
        b_mat: a `Tensor` of shape `[n, k]`.
        tol: (float) desired tolerance on the residual.
  
      Returns:
        X: `Tensor` `X` of shape `[n, k]` such that `A * X = B`.
  
      Raises:
        ValueError: if `a_mat` is not square or `a_mat` and `b_mat` have
        incompatible shapes.
      """
      _cg_check_shapes(a_mat, b_mat)
      n = tf.shape(b_mat)[0]
      x = tf.zeros_like(b_mat)
  
      r = b_mat - tf.matmul(a_mat, x)
      p = r
      rs_old = tf.einsum("ij,ij->j", r, r)
      rs_new = rs_old
  
      def body_fn(i, x, p, r, rs_old, rs_new):
          """One iteration of CG."""
          # Create a boolean mask shaped as [k] indicating the active columns, i.e.,
          # columns corresponding to large residuals. We only update variables
          # corresponding to those columns to avoid numerical issues.
          active_columns_mask = rs_old > tol
          # Replicate the mask along axis 0 to be of shape [n, k].
          active_columns_tiled_mask = tf.tile(
              tf.expand_dims(active_columns_mask, axis=0),
              multiples=[tf.shape(b_mat)[0], 1],
          )
          a_x_p = tf.matmul(a_mat, p)
          alpha = rs_old / tf.reduce_sum(p * a_x_p, axis=0)
          x = tf.where(
              active_columns_tiled_mask, x + tf.multiply(p, alpha), x
          )
          r = tf.where(
              active_columns_tiled_mask,
              r - tf.multiply(a_x_p, alpha),
              r,
          )
          rs_new = tf.where(
              active_columns_mask, tf.einsum("ij,ij->j", r, r), rs_new
          )
          p = tf.where(
              active_columns_tiled_mask,
              r + tf.multiply(p, rs_new / rs_old),
              p,
          )
          rs_old = tf.where(active_columns_mask, rs_new, rs_old)
          i = i + 1
          return i, x, p, r, rs_old, rs_new
  
      def while_exit_cond(i, x, p, r, rs_old, rs_new):
          """Exit the loop when n is reached or when the residual becomes small."""
          del x  # unused
          del p  # unused
          del r  # unused
          del rs_old  # unused
          i_cond = tf.less(i, n)
          residual_cond = tf.greater(
              tf.reduce_max(tf.sqrt(rs_new)), tol
          )
          return tf.logical_and(i_cond, residual_cond)
  
      _, x, _, _, _, _ = tf.while_loop(
          while_exit_cond,
          body_fn,
          [tf.constant(0), x, p, r, rs_old, rs_new],
          parallel_iterations=1,
      )
      return x
  
file_path: agents/tf_agents/bandits/policies/linalg.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 2 for calling a global function with example usage: conjugate_gradient(a_mat, b_mat, tol) and returns: x, i,  x,  p,  r,  rs_old,  rs_new, tf