jax-ml/jax · warning · NotImplementedError

Blocked version is not implemented yet.

Error message

Blocked version is not implemented yet.

What it means

jax.scipy.linalg.sqrtm implements only the point-wise (unblocked) Björck–Hammarling Schur method. Passing blocksize > 1 (SciPy-style parameter carried over in the JAX signature) raises NotImplementedError because the blocked algorithm has not been ported.

Source

Thrown at jax/_src/scipy/linalg.py:2223

     [0.92-0.71j 0.54-0.j   0.92+0.71j]]

    By definition, matrix multiplication of the matrix square root with itself should
    equal the input:

    >>> jnp.allclose(a, sqrt_a @ sqrt_a)
    Array(True, dtype=bool)

  Notes:
    This function implements the complex Schur method described in [1]_.  It does not use
    recursive blocking to speed up computations as a Sylvester Equation solver is not
    yet available in JAX.

  References:
    .. [1] Björck, Å., & Hammarling, S. (1983). "A Schur method for the square root of a matrix".
           Linear algebra and its applications, 52, 127-140.
  """
  if blocksize > 1:
      raise NotImplementedError("Blocked version is not implemented yet.")
  return _sqrtm(ensure_arraylike("scipy.sqrtm", A))


@jit
def _rsf2csf_2d(T: Array, Z: Array) -> tuple[Array, Array]:
  T, Z = promote_dtypes_complex(T, Z)
  eps = dtypes.finfo(T.dtype).eps
  N = T.shape[0]

  if N == 1:
    return T, Z

  def _update_T_Z(m, T, Z):
    mu = jnp_linalg.eigvals(lax.dynamic_slice(T, (m-1, m-1), (2, 2))) - T[m, m]
    r = jnp_linalg.norm(jnp.array([mu[0], T[m, m-1]])).astype(T.dtype)
    c = mu[0] / r
    s = T[m, m-1] / r
    G = jnp.array([[c.conj(), s], [-s, c]], dtype=T.dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the blocksize argument (default 1)
2. If blocking was a performance workaround, benchmark the JAX version first — the Schur method under jit is usually adequate
3. Use method='svd' sqrt via eigh on Hermitian positive-definite matrices if applicable

Example fix

# before
X = sqrtm(A, blocksize=64)
# after
X = sqrtm(A)  # blocksize defaults to 1

if blocksize != 1:
    blocksize = 1  # JAX only supports the unblocked algorithm
X = jax.scipy.linalg.sqrtm(A, blocksize=blocksize)

try:
    sqrtm(A, blocksize=blocksize)
except NotImplementedError:
    sqrtm(A)  # default blocksize=1

When it happens

Trigger: Calling sqrtm(A, blocksize=n) with n > 1; note blocksize=1 or the default (1) works fine.

Common situations: Copy-pasting SciPy call sites that tuned blocksize for performance; assuming the SciPy API surface is fully supported in JAX.

Related errors

• Only the type=1 case of eigh is implemented.
• Only the eigvals=None case of eigh is implemented.
• eigh_tridiagonal(..., select='v') is not implemented
• array ref with memory space only works inside of a `jit`.
• pinned array ref only works inside of a `jit`.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/1ec9851c133cd68c. Report an issue: GitHub.