jax-ml/jax · critical · ValueError

JAX requires ml_dtypes version 0.5 or newer; installed versi

Error message

JAX requires ml_dtypes version 0.5 or newer; installed version is {ml_dtypes.__version__}.

What it means

At import time, jax._src.dtypes verifies the installed ml_dtypes package is >= 0.5.0 because JAX relies on APIs introduced in that version (bfloat16/fp8 handling, etc.). An older version raises ValueError and JAX import fails.

Source

Thrown at jax/_src/dtypes.py:49

import ml_dtypes
import numpy as np

from jax._src import config
from jax._src import traceback_util
from jax._src.lib import _jax
from jax._src.typing import Array, DType, DTypeLike
from jax._src.util import StrictABC, set_module, cache

traceback_util.register_exclusion(__file__)

try:
  _ml_dtypes_version = tuple(map(int, ml_dtypes.__version__.split('.')[:3]))
except:
  pass
else:
  if _ml_dtypes_version < (0, 5):
    raise ValueError("JAX requires ml_dtypes version 0.5 or newer; "
                     f"installed version is {ml_dtypes.__version__}.")

export = set_module('jax.dtypes')

@export
class extended(np.generic):
  """Scalar class for extended dtypes.

  This is an abstract class that should never be instantiated, but rather
  exists for the sake of `jnp.issubdtype`.

  Examples:
    >>> from jax import random
    >>> from jax import dtypes
    >>> key = random.key(0)
    >>> jnp.issubdtype(key.dtype, dtypes.extended)
    True
  """

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. pip install -U 'ml_dtypes>=0.5'
2. Reinstall jax so pip resolves deps: pip install -U jax (or jax[cuda12]) which pulls a compatible ml_dtypes
3. Check for conflicting pins in requirements.txt/constraints and remove old ml_dtypes pins
4. In conda envs, update both jax and ml_dtypes from the same channel

Example fix

# before
ml_dtypes==0.4.0 in environment

# after
pip install -U 'ml_dtypes>=0.5'

import ml_dtypes
assert tuple(map(int, ml_dtypes.__version__.split('.')[:2])) >= (0, 5), 'upgrade: pip install -U ml_dtypes>=0.5'

When it happens

Trigger: Installing/upgrading jax while an old ml_dtypes (<0.5) remains in the environment; constraints files pinning ml_dtypes to an old version; stale conda environments.

Common situations: pip install jax without upgrading ml_dtypes; docker images with pinned old ml_dtypes; dependency resolvers keeping back ml_dtypes.

Related errors

• Requires libtpu >= 0.1.0
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32
• The names should be exclusive and should not intersect in `n

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