jax-ml/jax · error · std::invalid_argument

Unsupported data dtype: %s

Error message

Unsupported data dtype: %s

What it means

Maps NumPy dtypes to gpuDataType (cudaDataType/hipDataType) for GPU sparse/dense operations; unsupported dtypes (e.g. float64 on builds without it, uint types, bool) are rejected with the offending dtype's repr.

Source

Thrown at jaxlib/gpu/sparse.cc:73

      new absl::flat_hash_map<std::pair<char, int>, gpuDataType>({
          {{'f', 2}, GPU_R_16F},
          {{'c', 4}, GPU_C_16F},
          {{'f', 4}, GPU_R_32F},
          {{'c', 8}, GPU_C_32F},
          {{'f', 8}, GPU_R_64F},
          {{'c', 16}, GPU_C_64F},
#ifdef JAX_GPU_CUDA
          {{'i', 1}, CUDA_R_8I},
          {{'u', 1}, CUDA_R_8U},
          {{'i', 4}, CUDA_R_32I},
          {{'u', 4}, CUDA_R_32U},
          {{'V', 2}, CUDA_R_16BF},
#endif  // JAX_GPU_CUDA
      });
  auto it = types->find({np_type.kind(), np_type.itemsize()});
  if (it == types->end()) {
    nb::str repr = nb::repr(np_type);
    throw std::invalid_argument(
        absl::StrFormat("Unsupported data dtype: %s", repr.c_str()));
  }
  return it->second;
}
// Returns the descriptor for a Sparse matrix.
SparseMatDescriptor BuildSparseMatDescriptor(const dtype& data_dtype,
                                             const dtype& index_dtype, int rows,
                                             int cols, int nnz, int batch_count,
                                             int batch_stride) {
  gpuDataType value_type = DtypeToCudaDataType(data_dtype);
  gpusparseIndexType_t index_type = DtypeToCuSparseIndexType(index_dtype);
  return SparseMatDescriptor{value_type, index_type,  rows,        cols,
                             nnz,        batch_count, batch_stride};
}

// Returns the descriptor for a Dense matrix.
DenseMatDescriptor BuildDenseMatDescriptor(const dtype& data_dtype, int rows,
                                           int cols, int batch_count,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a supported dtype (float32, float64 where available, complex64/128, or bfloat16)
2. Cast inputs before calling the sparse op
3. Reinstall a jaxlib build matching your platform with full dtype support

Example fix

# before
spmm(a_bfloat16_unsupported_build, b)
# after
spmm(a.astype(jnp.float32), b)

SUPPORTED = {jnp.float32, jnp.float64, jnp.complex64, jnp.complex128}
assert x.dtype in SUPPORTED, f'unsupported {x.dtype}'

def supported_data_dtype(dt) -> bool:
    return dt in (jnp.float32, jnp.float64, jnp.complex64, jnp.complex128)

When it happens

Trigger: Calling jaxlib GPU sparse matmul descriptors with a data dtype unsupported by the compiled jaxlib GPU support table — e.g. complex, integer, or unsigned data in sparse ops, or bfloat16 on a CUDA build where it is compiled out.

Common situations: jax.experimentalsparse matmuls with unusual dtypes; ROCm/CUDA build differences where some dtype entries are ifdef'd out.

Related errors

• Sparse MMA unsupported for f32
• Unsupported dtype: {dtype}
• Expected metadata dtype to be uint2, got: {meta.dtype}
• Sparse metadata format not implemented for {operand_dtype=}
• cannot cast from `{dtype_name}`

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