{"record":{"id":"fee22a5013f4ede9","repo":"jax-ml/jax","slug":"a-dtype-a-dtype-is-not-supported","errorCode":null,"errorMessage":"A.dtype={A.dtype} is not supported.","messagePattern":"A\\.dtype=(.+?) is not supported\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/linalg.py","lineNumber":1472,"sourceCode":"  if A.dtype == 'float64' or A.dtype == 'complex128':\n   maxnorm = 5.371920351148152\n   n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))\n   A = A / 2 ** n_squarings.astype(A.dtype)\n   conds = jnp.array([1.495585217958292e-002, 2.539398330063230e-001,\n                      9.504178996162932e-001, 2.097847961257068e+000],\n                      dtype=A_L1.dtype)\n   idx = jnp.digitize(A_L1, conds)\n   U, V = lax.switch(idx, [_pade3, _pade5, _pade7, _pade9, _pade13], A)\n  elif A.dtype == 'float32' or A.dtype == 'complex64':\n    maxnorm = 3.925724783138660\n    n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))\n    A = A / 2 ** n_squarings.astype(A.dtype)\n    conds = jnp.array([4.258730016922831e-001, 1.880152677804762e+000],\n                      dtype=A_L1.dtype)\n    idx = jnp.digitize(A_L1, conds)\n    U, V = lax.switch(idx, [_pade3, _pade5, _pade7], A)\n  else:\n    raise TypeError(f\"A.dtype={A.dtype} is not supported.\")\n  P = U + V  # p_m(A) : numerator\n  Q = -U + V # q_m(A) : denominator\n  return P, Q, n_squarings\n\ndef _solve_P_Q(P: ArrayLike, Q: ArrayLike, upper_triangular: bool = False) -> Array:\n  if upper_triangular:\n    return solve_triangular(Q, P)\n  else:\n    return jnp_linalg.solve(Q, P)\n\ndef _precise_dot(A: ArrayLike, B: ArrayLike) -> Array:\n  return jnp.dot(A, B, precision=lax.Precision.HIGHEST)\n\n@jit(static_argnums=2)\ndef _squaring(R: Array, n_squarings: Array, max_squarings: int) -> Array:\n  # squaring step to undo scaling\n  def _squaring_precise(x):\n    return _precise_dot(x, x)","sourceCodeStart":1454,"sourceCodeEnd":1490,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/linalg.py#L1454-L1490","documentation":"expm's Padé-squaring implementation only supports float32/complex64 and float64/complex128 inputs. _calc_P_Q selects norm thresholds per precision and raises TypeError for any other dtype, including bfloat16, float16, and integer inputs.","triggerScenarios":"Calling jax.scipy.linalg.expm on a bfloat16 or float16 array (common after promotion in mixed-precision models or TPU defaults), or on an integer matrix.","commonSituations":"Mixed-precision training on TPU where arrays default to bfloat16; passing an un-promoted int array.","solutions":["Cast to float64 or complex128: A = A.astype(jnp.float64)","Call expm with at least float32: use jax.scipy.linalg.expm(A.astype(jnp.float32))","Enable x64 if you need double precision: jax.config.update('jax_enable_x64', True) before casting to float64"],"exampleFix":"// before\nM = jax.scipy.linalg.expm(A_bf16)\n// after\nM = jax.scipy.linalg.expm(A_bf16.astype(jnp.float32))","handlingStrategy":"type-guard","validationCode":"if A.dtype not in (jnp.float32, jnp.float64, jnp.complex64, jnp.complex128): A = A.astype(jnp.float32)","typeGuard":"_EXPM_DTYPES = {'float32','float64','complex64','complex128'}\ndef expm_dtype_ok(A): return str(A.dtype) in _EXPM_DTYPES","tryCatchPattern":null,"preventionTips":["Cast to float32/float64 before expm in mixed-precision pipelines","Enable jax_enable_x64 when double precision matters"],"tags":["jax","linalg","dtype","matrix-exponential"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}