{"record":{"id":"9dc44d3df7c26adb","repo":"jax-ml/jax","slug":"unsupported-dtype-dtype","errorCode":null,"errorMessage":"Unsupported dtype: {dtype}","messagePattern":"Unsupported dtype: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/linalg.py","lineNumber":1098,"sourceCode":"\n\ndef _eig_gpu_lowering(ctx, operand, *,\n                      compute_left_eigenvectors, compute_right_eigenvectors,\n                      enable_eigvec_derivs, implementation, target_name_prefix):\n  del enable_eigvec_derivs\n  operand_aval, = ctx.avals_in\n  batch_dims = operand_aval.shape[:-2]\n  n, m = operand_aval.shape[-2:]\n  assert n == m\n\n  dtype = operand_aval.dtype\n  complex_dtype = np.result_type(dtype, 1j)\n  if dtype in (np.float32, np.float64):\n    is_real = True\n  elif dtype in (np.complex64, np.complex128):\n    is_real = False\n  else:\n    raise ValueError(f\"Unsupported dtype: {dtype}\")\n\n  have_cusolver_geev = (\n      target_name_prefix == \"cu\"\n      and cuda_versions\n      and cuda_versions.cusolver_get_version() >= 11701\n  )\n\n  if (\n      implementation is None and have_cusolver_geev\n      and not compute_left_eigenvectors\n  ) or implementation == EigImplementation.CUSOLVER:\n    if not have_cusolver_geev:\n      raise RuntimeError(\n          \"Nonsymmetric eigendecomposition requires cusolver 11.7.1 or newer\"\n      )\n    if compute_left_eigenvectors:\n      raise NotImplementedError(\n          \"Left eigenvectors are not supported by cusolver\")","sourceCodeStart":1080,"sourceCodeEnd":1116,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/linalg.py#L1080-L1116","documentation":"jax/_src/lax/linalg.py:1098 in _eig_gpu_lowering (eig GPU path). Thrown when the dtype of the matrix passed to jax.lax.linalg.eig is not one of the four supported: float32, float64, complex64, complex128. Before lowering to a GPU kernel, JAX classifies the input as real or complex; anything else (e.g. int, float16/bfloat16, or an extended-precision type) reaches the else branch.","triggerScenarios":"Calling jax.lax.linalg.eig (or jax.numpy.linalg.eig) on GPU with an operand dtype outside {float32, float64, complex64, complex128}, e.g. an integer array, float16/bfloat16 weights, or a dtype created via custom casting. Only the GPU lowering raises; CPU may fail elsewhere or differently.","commonSituations":"Models storing weights in bfloat16 (common in transformers) passed directly to eig; integer matrices from preprocessing; accidental object/odd dtype after mixing NumPy and JAX arrays.","solutions":["Cast the operand to a supported dtype before calling eig: a.astype(jnp.float32) (or complex64/float64/complex128 as appropriate)","Check and normalize dtype at model input boundaries, e.g. jax.tree_util.map over params enforcing f32/f64","If you truly need float16/bfloat16 results, compute in float32 and cast the outputs back down"],"exampleFix":"// before\nw, v = jax.lax.linalg.eig(a)  # a is bfloat16\n// after\nw, v = jax.lax.linalg.eig(a.astype(jnp.float32))","handlingStrategy":"validation","validationCode":"SUPPORTED = {jnp.float32, jnp.float64, np.complex64, np.complex128}\nif a.dtype not in SUPPORTED:\n    a = a.astype(jnp.float32)","typeGuard":"def is_eig_dtype(a: jax.Array) -> bool:\n    return a.dtype in (jnp.float32, jnp.float64, jnp.complex64, jnp.complex128)","tryCatchPattern":null,"preventionTips":["Normalize dtypes at input boundaries with jnp.asarray(x, jnp.float32)","Watch for bfloat16 params from mixed-precision training before linalg ops"],"tags":["jax","linalg","eig","dtype","gpu"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}