# Source code for jax._src.third_party.scipy.linalg

```from typing import Callable, Tuple

import scipy.linalg

from jax import jit, lax
from jax._src.numpy import lax_numpy as jnp
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

@jit
def _algorithm_11_1_1(F: Array, T: Array) -> Tuple[Array, Array]:
# Algorithm 11.1.1 from Golub and Van Loan "Matrix Computations"
N = T.shape[0]
minden = jnp.abs(T[0, 0])

def _outer_loop(p, F_minden):
_, F, minden = lax.fori_loop(1, N-p+1, _inner_loop, (p, *F_minden))
return F, minden

def _inner_loop(i, p_F_minden):
p, F, minden = p_F_minden
j = i+p
s = T[i-1, j-1] * (F[j-1, j-1] - F[i-1, i-1])
T_row, T_col = T[i-1], T[:, j-1]
F_row, F_col = F[i-1], F[:, j-1]
ind = (jnp.arange(N) >= i) & (jnp.arange(N) < j-1)
val = (jnp.where(ind, T_row, 0) @ jnp.where(ind, F_col, 0) -
jnp.where(ind, F_row, 0) @ jnp.where(ind, T_col, 0))
s = s + val
den = T[j-1, j-1] - T[i-1, i-1]
s = jnp.where(den != 0, s / den, s)
F = F.at[i-1, j-1].set(s)
minden = jnp.minimum(minden, jnp.abs(den))
return p, F, minden

return lax.fori_loop(1, N, _outer_loop, (F, minden))

_FUNM_LAX_DESCRIPTION = """\
The array returned by :py:func:`jax.scipy.linalg.funm` may differ in dtype
from the array returned by py:func:`scipy.linalg.funm`. Specifically, in cases
where all imaginary parts of the array values are close to zero, the SciPy
function may return a real-valued array, whereas the JAX implementation will
return a complex-valued array.

Additionally, unlike the SciPy implementation, when ``disp=True`` no warning
will be printed if the error in the array output is estimated to be large.
"""

[docs]@_wraps(scipy.linalg.funm, lax_description=_FUNM_LAX_DESCRIPTION)
def funm(A: ArrayLike, func: Callable[[Array], Array], disp: bool = True) -> Tuple[Array, Array]:
A = jnp.asarray(A)
if A.ndim != 2 or A.shape[0] != A.shape[1]:
raise ValueError('expected square array_like input')

T, Z = schur(A)
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
```