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_kernelparams(self):
      skip_if_no_xla(self.skipTest)
      if not tf.executing_eagerly():
          return  # jit_compile is eager-only.
      feature_dim = 3
  
      def kernel_fn(eq_params, poly_params):
          return exponentiated_quadratic.ExponentiatedQuadratic(
              *eq_params
          ) * polynomial.Polynomial(
              bias_amplitude=poly_params[0], shift=poly_params[1]
          )
  
      # TODO(b/284106340): Return this to a dictionary.
      kernel_args = (
          (
              tf.random.uniform(
                  [], 1.5, 2.5, dtype=tf.float64
              ),  # amplitude
              tf.random.uniform([], 0.5, 1.5, dtype=tf.float64),
          ),  # length_scale
          (
              tf.random.uniform(
                  [feature_dim],
                  0.5,
                  1.5,  # bias_amplitude
                  dtype=tf.float64,
              ),
              tf.random.normal([feature_dim], dtype=tf.float64),
          ),
      )  # shift
  
      x1 = tf.random.normal([5, feature_dim], dtype=tf.float64)
      x2 = tf.random.normal([7, feature_dim], dtype=tf.float64)
      linop = linear_operator_psd_kernel.LinearOperatorPSDKernel(
          kernel_fn, x1, x2, kernel_args=kernel_args, num_matmul_parts=3
      )
      x = tf.random.normal([7, 2], dtype=tf.float64)
  
      @tf.function(jit_compile=True)
      def f():
          with tf.GradientTape() as tape:
              tape.watch((x1, x2, x, kernel_args))
              y = linop.matmul(x)
  
          out_grad = tf.random.normal(tf.shape(y), dtype=tf.float64)
          actuals = tape.gradient(
              y, (x1, x2, x, kernel_args), output_gradients=out_grad
          )
          return y, actuals, out_grad
  
      y, actuals, out_grad = f()
  
      with tf.GradientTape() as tape:
          tape.watch((x1, x2, x, kernel_args))
          y = tf.matmul(kernel_fn(*kernel_args).matrix(x1, x2), x)
      expecteds = tape.gradient(
          y, (x1, x2, x, kernel_args), output_gradients=out_grad
      )
  
      expecteds, actuals = self.evaluate([expecteds, actuals])
  
      tf.nest.assert_same_structure(expecteds, actuals)
      for expected, actual in zip(
          tf.nest.flatten(expecteds), tf.nest.flatten(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 14 for calling the LinearOperatorPSDKernelTest obj's underlying member method code with example usage: obj.test_matmul_grad_xla_kernelparams() and returns: unknown, y,  actuals,  out_grad