code: def testConjugateGradientMultipleRHS(self, n, rhs):
      x_obs = tf.constant(
          np.random.rand(n, 2), dtype=tf.float32, shape=[n, 2]
      )
      a_mat = tf.eye(n) + tf.matmul(
          x_obs, tf.linalg.matrix_transpose(x_obs)
      )
      x_exact = tf.constant(
          np.random.rand(n, rhs), dtype=tf.float32, shape=[n, rhs]
      )
      b_mat = tf.matmul(a_mat, x_exact)
      x_approx = self.evaluate(linalg.conjugate_gradient(a_mat, b_mat))
      x_exact_numpy = self.evaluate(x_exact)
      self.assertAllClose(x_exact_numpy, x_approx, rtol=1e-4, atol=1e-4)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.policies import linalg
label_desc: Use image node_id 2 for calling the ConjugateGradientTest obj's underlying member method code with example usage: obj.testConjugateGradientMultipleRHS(n, rhs) without return types