code: def _check_inputs(A, X):
      n, k = X.shape
      dt = X.dtype
  
      if k == 0:
          raise ValueError(f"must have search dim > 0, got {k}")
  
      if k * 5 >= n:
          raise ValueError(
              f"expected search dim * 5 < matrix dim (got {k * 5}, {n})"
          )
  
      test_output = A(jnp.zeros((n, 1), dtype=X.dtype))
  
      if test_output.dtype != dt:
          raise ValueError(
              f"A, X must have same dtypes (were {test_output.dtype}, {dt})"
          )
  
      if test_output.shape != (n, 1):
          s = test_output.shape
          raise ValueError(
              f"A must be ({n}, {n}) matrix A, got output {s}"
          )
  
imports: from __future__ import annotations
  from typing import Union, Callable
  import functools
  import jax
  import jax.numpy
  from jax.experimental import sparse
  from jax.interpreters import mlir
  from jax.interpreters import xla
  from jax._src import core
  from jax._src.interpreters import ad
  from jax._src.lib import gpu_solver
  import numpy
  from scipy.sparse import csr_matrix, linalg
label_desc: Use image node_id 4 for calling a global function with example usage: _check_inputs(A, X) without return types