jax-ml/jax · error · TypeError

A.dtype={A.dtype} is not supported.

Error message

A.dtype={A.dtype} is not supported.

What it means

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.

Source

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

  if A.dtype == 'float64' or A.dtype == 'complex128':
   maxnorm = 5.371920351148152
   n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
   A = A / 2 ** n_squarings.astype(A.dtype)
   conds = jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
                      9.504178996162932e-001, 2.097847961257068e+000],
                      dtype=A_L1.dtype)
   idx = jnp.digitize(A_L1, conds)
   U, V = lax.switch(idx, [_pade3, _pade5, _pade7, _pade9, _pade13], A)
  elif A.dtype == 'float32' or A.dtype == 'complex64':
    maxnorm = 3.925724783138660
    n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
    A = A / 2 ** n_squarings.astype(A.dtype)
    conds = jnp.array([4.258730016922831e-001, 1.880152677804762e+000],
                      dtype=A_L1.dtype)
    idx = jnp.digitize(A_L1, conds)
    U, V = lax.switch(idx, [_pade3, _pade5, _pade7], A)
  else:
    raise TypeError(f"A.dtype={A.dtype} is not supported.")
  P = U + V  # p_m(A) : numerator
  Q = -U + V # q_m(A) : denominator
  return P, Q, n_squarings

def _solve_P_Q(P: ArrayLike, Q: ArrayLike, upper_triangular: bool = False) -> Array:
  if upper_triangular:
    return solve_triangular(Q, P)
  else:
    return jnp_linalg.solve(Q, P)

def _precise_dot(A: ArrayLike, B: ArrayLike) -> Array:
  return jnp.dot(A, B, precision=lax.Precision.HIGHEST)

@jit(static_argnums=2)
def _squaring(R: Array, n_squarings: Array, max_squarings: int) -> Array:
  # squaring step to undo scaling
  def _squaring_precise(x):
    return _precise_dot(x, x)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast to float64 or complex128: A = A.astype(jnp.float64)
2. Call expm with at least float32: use jax.scipy.linalg.expm(A.astype(jnp.float32))
3. Enable x64 if you need double precision: jax.config.update('jax_enable_x64', True) before casting to float64

Example fix

// before
M = jax.scipy.linalg.expm(A_bf16)
// after
M = jax.scipy.linalg.expm(A_bf16.astype(jnp.float32))

if A.dtype not in (jnp.float32, jnp.float64, jnp.complex64, jnp.complex128): A = A.astype(jnp.float32)

_EXPM_DTYPES = {'float32','float64','complex64','complex128'}
def expm_dtype_ok(A): return str(A.dtype) in _EXPM_DTYPES

When it happens

Trigger: 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.

Common situations: Mixed-precision training on TPU where arrays default to bfloat16; passing an un-promoted int array.

Related errors

• Unsupported dtype: {dtype}
• Expected A to be a (batched) square matrix, got {A.shape=}.
• diagonal and off-diagonal values must have same dtype, got {
• Only float32 and float64 inputs are supported as inputs to j
• dtype must be a complex floating-point type; got {dtype}.

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