{"record":{"id":"1ec9851c133cd68c","repo":"jax-ml/jax","slug":"blocked-version-is-not-implemented-yet","errorCode":null,"errorMessage":"Blocked version is not implemented yet.","messagePattern":"Blocked version is not implemented yet\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"warning","filePath":"jax/_src/scipy/linalg.py","lineNumber":2223,"sourceCode":"     [0.92-0.71j 0.54-0.j   0.92+0.71j]]\n\n    By definition, matrix multiplication of the matrix square root with itself should\n    equal the input:\n\n    >>> jnp.allclose(a, sqrt_a @ sqrt_a)\n    Array(True, dtype=bool)\n\n  Notes:\n    This function implements the complex Schur method described in [1]_.  It does not use\n    recursive blocking to speed up computations as a Sylvester Equation solver is not\n    yet available in JAX.\n\n  References:\n    .. [1] Björck, Å., & Hammarling, S. (1983). \"A Schur method for the square root of a matrix\".\n           Linear algebra and its applications, 52, 127-140.\n  \"\"\"\n  if blocksize > 1:\n      raise NotImplementedError(\"Blocked version is not implemented yet.\")\n  return _sqrtm(ensure_arraylike(\"scipy.sqrtm\", A))\n\n\n@jit\ndef _rsf2csf_2d(T: Array, Z: Array) -> tuple[Array, Array]:\n  T, Z = promote_dtypes_complex(T, Z)\n  eps = dtypes.finfo(T.dtype).eps\n  N = T.shape[0]\n\n  if N == 1:\n    return T, Z\n\n  def _update_T_Z(m, T, Z):\n    mu = jnp_linalg.eigvals(lax.dynamic_slice(T, (m-1, m-1), (2, 2))) - T[m, m]\n    r = jnp_linalg.norm(jnp.array([mu[0], T[m, m-1]])).astype(T.dtype)\n    c = mu[0] / r\n    s = T[m, m-1] / r\n    G = jnp.array([[c.conj(), s], [-s, c]], dtype=T.dtype)","sourceCodeStart":2205,"sourceCodeEnd":2241,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/linalg.py#L2205-L2241","documentation":"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.","triggerScenarios":"Calling sqrtm(A, blocksize=n) with n > 1; note blocksize=1 or the default (1) works fine.","commonSituations":"Copy-pasting SciPy call sites that tuned blocksize for performance; assuming the SciPy API surface is fully supported in JAX.","solutions":["Drop the blocksize argument (default 1)","If blocking was a performance workaround, benchmark the JAX version first — the Schur method under jit is usually adequate","Use method='svd' sqrt via eigh on Hermitian positive-definite matrices if applicable"],"exampleFix":"# before\nX = sqrtm(A, blocksize=64)\n# after\nX = sqrtm(A)  # blocksize defaults to 1","handlingStrategy":"validation","validationCode":"if blocksize != 1:\n    blocksize = 1  # JAX only supports the unblocked algorithm\nX = jax.scipy.linalg.sqrtm(A, blocksize=blocksize)","typeGuard":null,"tryCatchPattern":"try:\n    sqrtm(A, blocksize=blocksize)\nexcept NotImplementedError:\n    sqrtm(A)  # default blocksize=1","preventionTips":["Don't carry SciPy-only performance knobs into JAX call sites","Diff the JAX function signature against scipy's before porting","Treat NotImplementedError as a version/feature probe, then fall back"],"tags":["jax","sqrtm","not-implemented","linalg"],"backgroundTag":"unsupported-operation-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}