class_docstr: Tests for linear_operator_psd_kernel.LinearOperatorPSDKernel.
class_docstr_tok: ["Tests","for","linear_operator_psd_kernel.LinearOperatorPSDKernel","."]
code: def test_col_batch(self):
      kernel = exponentiated_quadratic.ExponentiatedQuadratic()
      x1 = tf.random.normal([3, 5, 2])
      x2 = tf.random.normal([7, 1, 4, 2])
      linop = linear_operator_psd_kernel.LinearOperatorPSDKernel(
          kernel, x1, x2
      )
      j = np.random.randint(0, 4, size=(7, 3))
      cov = kernel.matrix(x1, x2)
      transpose = tf.linalg.matrix_transpose
      # Gather with batch_dims wants all the batch dims adjacent and leading, so
      # transpose-gather-transpose is easier to write than injecting a
      # range(nrows) column into the gather indices.
      expected, actual = self.evaluate(
          [
              transpose(
                  tf.gather(
                      transpose(cov), j[..., tf.newaxis], batch_dims=2
                  )
              )[..., 0],
              linop.col(j),
          ]
      )
      self.assertAllClose(expected, actual)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.linalg import linear_operator_psd_kernel
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math.psd_kernels import exponentiated_quadratic
  from tensorflow_probability.python.math.psd_kernels import polynomial
label_desc: Use image node_id 7 for calling the LinearOperatorPSDKernelTest obj's underlying member method code with example usage: obj.test_col_batch() without return types