code: def _gradientTest(self, diags, rhs, dtype=dtypes.float64):
      def reference_matmul(diags, rhs):
          matrix = self._makeTridiagonalMatrix(
              diags[..., 0, :-1], diags[..., 1, :], diags[..., 2, 1:]
          )
          return math_ops.matmul(matrix, rhs)
  
      with self.session(), self.test_scope():
          diags = constant_op.constant(diags, dtype=dtype)
          rhs = constant_op.constant(rhs, dtype=dtype)
          grad_reference, _ = gradient_checker_v2.compute_gradient(
              reference_matmul, [diags, rhs]
          )
          (
              grad_theoretical,
              grad_numerical,
          ) = gradient_checker_v2.compute_gradient(
              linalg_impl.tridiagonal_matmul, [diags, rhs]
          )
      self.assertAllClose(grad_theoretical, grad_numerical)
      self.assertAllClose(grad_theoretical, grad_reference)
  
imports: import numpy
  from tensorflow.compiler.tests import xla_test
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import gradient_checker_v2
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.ops.linalg import linalg_impl
  from tensorflow.python.platform import googletest
label_desc: Use image node_id 5 for calling the TridiagonalMatMulOpsTest obj's underlying member method code with example usage: obj._gradientTest(diags, rhs, dtype) and returns: math_ops