code: def funm(
      A: ArrayLike, func: Callable[[Array], Array], disp: bool = True
  ) -> Array | tuple[Array, Array]:
      A_arr = jnp.asarray(A)
      if A_arr.ndim != 2 or A_arr.shape[0] != A_arr.shape[1]:
          raise ValueError("expected square array_like input")
  
      T, Z = schur(A_arr)
      T, Z = rsf2csf(T, Z)
  
      F = jnp.diag(func(jnp.diag(T)))
      F = F.astype(T.dtype.char)
  
      F, minden = _algorithm_11_1_1(F, T)
      F = Z @ F @ Z.conj().T
  
      if disp:
          return F
  
      if F.dtype.char.lower() == "e":
          tol = jnp.finfo(jnp.float16).eps
      if F.dtype.char.lower() == "f":
          tol = jnp.finfo(jnp.float32).eps
      else:
          tol = jnp.finfo(jnp.float64).eps
  
      minden = jnp.where(minden == 0.0, tol, minden)
      err = jnp.where(
          jnp.any(jnp.isinf(F)),
          jnp.inf,
          jnp.minimum(
              1,
              jnp.maximum(
                  tol, (tol / minden) * norm(jnp.triu(T, 1), 1)
              ),
          ),
      )
  
      return F, err
  
file_path: jax/jax/_src/third_party/scipy/linalg.py
imports: from __future__ import annotations
  from typing import Callable
  import scipy.linalg
  from jax import jit, lax
  import jax.numpy
  from jax._src.numpy.linalg import norm
  from jax._src.numpy.util import _wraps
  from jax._src.scipy.linalg import rsf2csf, schur
  from jax._src.typing import ArrayLike, Array
label_desc: Use image node_id 2 for calling a global function with example usage: funm(A, func, disp) and returns: F,  err, F