code: def _normalize_matvec(f):
      """Normalize an argument for computing matrix-vector products."""
      if callable(f):
          return f
      elif isinstance(f, (np.ndarray, jax.Array)):
          if f.ndim != 2 or f.shape[0] != f.shape[1]:
              raise ValueError(
                  f"linear operator must be a square matrix, but has shape: {f.shape}"
              )
          return partial(_dot, f)
      elif hasattr(f, "__matmul__"):
          if (
              hasattr(f, "shape")
              and len(f.shape) != 2
              or f.shape[0] != f.shape[1]
          ):
              raise ValueError(
                  f"linear operator must be a square matrix, but has shape: {f.shape}"
              )
          return partial(operator.matmul, f)
      else:
          raise TypeError(
              f"linear operator must be either a function or ndarray: {f}"
          )
  
imports: from functools import partial
  import operator
  import numpy
  import jax
  import jax.numpy
  from jax import device_put
  from jax import lax
  from jax import scipy
  from jax.tree_util import tree_leaves, tree_map, tree_structure, tree_reduce, Partial
  from jax._src import dtypes
  from jax._src.lax import lax
  from jax._src.util import safe_map
label_desc: Use image node_id 8 for calling a global function with example usage: _normalize_matvec(f) and returns: f, partial, partial