code: def _testAllFormats(
      self, superdiag, maindiag, subdiag, rhs, dtype=dtypes.float64
  ):
      superdiag_extended = np.pad(superdiag, [0, 1], "constant")
      subdiag_extended = np.pad(subdiag, [1, 0], "constant")
  
      diags_compact = np.stack(
          [superdiag_extended, maindiag, subdiag_extended]
      )
      diags_matrix = (
          np.diag(superdiag, 1)
          + np.diag(maindiag, 0)
          + np.diag(subdiag, -1)
      )
  
      with self.test_scope():
          diags_sequence_op = (
              constant_op.constant(superdiag_extended, dtype),
              constant_op.constant(maindiag, dtype),
              constant_op.constant(subdiag_extended, dtype),
          )
          diags_compact_op = constant_op.constant(diags_compact, dtype)
          diags_matrix_op = constant_op.constant(diags_matrix, dtype)
          rhs_op = constant_op.constant(rhs, dtype)
  
          rhs_batch_op = array_ops_stack.stack([rhs_op, 2 * rhs_op])
          diags_compact_batch_op = array_ops_stack.stack(
              [diags_compact_op, 2 * diags_compact_op]
          )
          diags_matrix_batch_op = array_ops_stack.stack(
              [diags_matrix_op, 2 * diags_matrix_op]
          )
          diags_sequence_batch_op = [
              array_ops_stack.stack([x, 2 * x])
              for x in diags_sequence_op
          ]
  
          results = [
              self._jit_tridiagonal_matmul(
                  diags_sequence_op, rhs_op, diagonals_format="sequence"
              ),
              self._jit_tridiagonal_matmul(
                  diags_compact_op, rhs_op, diagonals_format="compact"
              ),
              self._jit_tridiagonal_matmul(
                  diags_matrix_op, rhs_op, diagonals_format="matrix"
              ),
          ]
  
          results_batch = [
              self._jit_tridiagonal_matmul(
                  diags_sequence_batch_op,
                  rhs_batch_op,
                  diagonals_format="sequence",
              ),
              self._jit_tridiagonal_matmul(
                  diags_compact_batch_op,
                  rhs_batch_op,
                  diagonals_format="compact",
              ),
              self._jit_tridiagonal_matmul(
                  diags_matrix_batch_op,
                  rhs_batch_op,
                  diagonals_format="matrix",
              ),
          ]
  
          expected = self._tridiagonal_matmul(
              diags_sequence_op, rhs_op, diagonals_format="sequence"
          )
          expected_batch = np.stack([expected, 4 * expected])
          for result in results:
              self.assertAllClose(result, expected)
          for result in results_batch:
              self.assertAllClose(result, expected_batch)
  
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 3 for calling the TridiagonalMatMulOpsTest obj's underlying member method code with example usage: obj._testAllFormats(superdiag, maindiag, subdiag, rhs, dtype) without return types