class_docstr: Tests for linear_operator_psd_kernel.LinearOperatorPSDKernel.
class_docstr_tok: ["Tests","for","linear_operator_psd_kernel.LinearOperatorPSDKernel","."]
code: def test_matmul_grad_xla(self):
      skip_if_no_xla(self.skipTest)
      if not tf.executing_eagerly():
          return  # jit_compile is eager-only.
      kernel = exponentiated_quadratic.ExponentiatedQuadratic()
      x1 = tf.random.normal([5, 3])
      x2 = tf.random.normal([7, 3])
      linop = linear_operator_psd_kernel.LinearOperatorPSDKernel(
          kernel, x1, x2, num_matmul_parts=3
      )
      x = tf.random.normal([7, 2])
  
      @tf.function(jit_compile=True)
      def f():
          with tf.GradientTape() as tape:
              tape.watch((x1, x2, x))
              y = linop.matmul(x)
  
          out_grad = tf.random.normal(tf.shape(y))
          actuals = tape.gradient(
              y, (x1, x2, x), output_gradients=out_grad
          )
          return y, actuals, out_grad
  
      y, actuals, out_grad = f()
  
      with tf.GradientTape() as tape:
          tape.watch((x1, x2, x))
          y = tf.matmul(kernel.matrix(x1, x2), x)
      expecteds = tape.gradient(
          y, (x1, x2, x), output_gradients=out_grad
      )
  
      expecteds, actuals = self.evaluate([expecteds, actuals])
  
      self.assertEqual(len(expecteds), len(actuals))
      for expected, actual in zip(expecteds, actuals):
          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 13 for calling the LinearOperatorPSDKernelTest obj's underlying member method code with example usage: obj.test_matmul_grad_xla() and returns: y,  actuals,  out_grad