jax-ml/jax · error · TypeError

Expected a string or dtype-like object; got {dtype=}

Error message

Expected a string or dtype-like object; got {dtype=}

What it means

jax.dlpack.is_supported_dtype rejects dtype=None because NumPy would silently interpret None as float64, producing a surprising dtype instead of an error. JAX raises an explicit TypeError instead.

Source

Thrown at jax/_src/dlpack.py:57

# For example,
# hash(jnp.float32) != hash(jnp.dtype(jnp.float32))
# hash(jnp.float32) == hash(jnp.dtype(jnp.float32).type)

# TODO(vanderplas): remove this set
SUPPORTED_DTYPES: frozenset[DTypeLike] = frozenset({
    jnp_types.int8, jnp_types.int16, jnp_types.int32, jnp_types.int64,
    jnp_types.uint8, jnp_types.uint16, jnp_types.uint32, jnp_types.uint64,
    jnp_types.float16, jnp_types.bfloat16, jnp_types.float32, jnp_types.float64,
    jnp_types.complex64, jnp_types.complex128, jnp_types.bool_})

SUPPORTED_DTYPES_SET: frozenset[np.dtype] = frozenset({np.dtype(dt) for dt in SUPPORTED_DTYPES})


def is_supported_dtype(dtype: DTypeLike) -> bool:
  """Check if dtype is supported by jax.dlpack."""
  if dtype is None:
    # NumPy will silently cast this to float64, which may be surprising.
    raise TypeError(f"Expected a string or dtype-like object; got {dtype=}")
  return np.dtype(dtype) in SUPPORTED_DTYPES_SET


def _to_dlpack(x: Array, stream: int | Any | None,
               src_device: _jax.Device | None = None,
               device: _jax.Device | None = None,
               copy: bool | None = None):

  if src_device is None:
    src_device, = x.devices()
  if device and (src_device is None or device != src_device):
    if copy is not None and not copy:
      raise ValueError(
        f"Specified {device=} which requires a copy since the source device "
        f"is {repr(src_device)}, however copy=False. Set copy=True or "
        "copy=None to perform the requested operation."
      )
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Default the dtype explicitly (e.g. `dtype or torch.float32`) before calling
2. Validate dtype is not None with an isinstance/type check upstream
3. Pass a concrete numpy/torch dtype string such as 'float32'

Example fix

# before
jax.dlpack.is_supported_dtype(cfg.dtype)  # cfg.dtype is None

# after
dtype = cfg.dtype or 'float32'
jax.dlpack.is_supported_dtype(dtype)

if dtype is None:
    raise ValueError('dtype must be specified for dlpack conversion')
jax.dlpack.is_supported_dtype(dtype)

def has_dtype(d) -> bool:
    return d is not None and not isinstance(d, type(None))

try:
    jax.dlpack.is_supported_dtype(dtype)
except TypeError:
    dtype = 'float32'

When it happens

Trigger: Calling jax.dlpack.is_supported_dtype(None), or passing an unpopulated dtype variable (common when a dtype field of a config/dataclass defaults to None) into dlpack conversion.

Common situations: Optional dtype parameters threaded from user config into from_dlpack/to_dlpack paths; passing torch tensors' dtype when it is None.

Related errors

• dtype cannot be None.
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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