class_docstr: Tests for linear_operator_psd_kernel.LinearOperatorPSDKernel.
class_docstr_tok: ["Tests","for","linear_operator_psd_kernel.LinearOperatorPSDKernel","."]
code: def test_diag_part(self):
      kernel = exponentiated_quadratic.ExponentiatedQuadratic()
      x1 = tf.random.normal([7, 3, 5, 2])  # square matrix 5x5
      linop = linear_operator_psd_kernel.LinearOperatorPSDKernel(
          kernel, x1
      )
      expected, actual = self.evaluate(
          [
              tf.linalg.diag_part(kernel.matrix(x1, x1)),
              linop.diag_part(),
          ]
      )
      self.assertAllClose(expected, actual)
  
      x2 = tf.random.normal([3, 11, 2])  # wide matrix 5x11
      linop = linear_operator_psd_kernel.LinearOperatorPSDKernel(
          kernel, x1, x2
      )
      expected, actual = self.evaluate(
          [
              tf.linalg.diag_part(kernel.matrix(x1, x2)),
              linop.diag_part(),
          ]
      )
      self.assertAllClose(expected, actual)
  
      x2 = tf.random.normal([2, 2])  # tall matrix 5x2
      linop = linear_operator_psd_kernel.LinearOperatorPSDKernel(
          kernel, x1, x2
      )
      expected, actual = self.evaluate(
          [
              tf.linalg.diag_part(kernel.matrix(x1, x2)),
              linop.diag_part(),
          ]
      )
      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 2 for calling the LinearOperatorPSDKernelTest obj's underlying member method code with example usage: obj.test_diag_part() without return types