docstr: The Arnoldi decomposition within GMRES is correct.
docstr_tok: ["The","Arnoldi","decomposition","within","GMRES","is","correct","."]
code: def test_gmres_arnoldi_step(self, shape, dtype, preconditioner):
      """
      The Arnoldi decomposition within GMRES is correct.
      """
      if not config.enable_x64.value:
          raise unittest.SkipTest("requires x64 mode")
  
      rng = jtu.rand_default(self.rng())
      A = rng(shape, dtype)
      M = self._fetch_preconditioner(preconditioner, A, rng=rng)
      if preconditioner is None:
          M = lambda x: x
      else:
          M = partial(matmul_high_precision, M)
      n = shape[0]
      x0 = rng(shape[:1], dtype)
      Q = np.zeros((n, n + 1), dtype=dtype)
      Q[:, 0] = x0 / jnp.linalg.norm(x0).astype(dtype)
      Q = jnp.array(Q)
      H = jnp.eye(n, n + 1, dtype=dtype)
  
      @jax.tree_util.Partial
      def A_mv(x):
          return matmul_high_precision(A, x)
  
      for k in range(n):
          Q, H, _ = sp_linalg._kth_arnoldi_iteration(k, A_mv, M, Q, H)
      QA = matmul_high_precision(Q[:, :n].conj().T, A)
      QAQ = matmul_high_precision(QA, Q[:, :n])
      self.assertAllClose(QAQ, H.T[:n, :], rtol=2e-5, atol=2e-5)
  
imports: from functools import partial
  import unittest
  from absl.testing import absltest
  import numpy
  import scipy.sparse.linalg
  import jax
  import jax.numpy
  import jax.scipy.sparse.linalg
  from jax import jit
  from jax import lax
  from jax.tree_util import register_pytree_node_class
  import jax._src.scipy.sparse.linalg
  from jax._src import config
  from jax._src import dtypes
  from jax._src import test_util
label_desc: Use image node_id 21 for calling the LaxBackedScipyTests obj's underlying member method code with example usage: obj.test_gmres_arnoldi_step(shape, dtype, preconditioner) and returns: matmul_high_precision