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

Unsupported index dtype: %s

Error message

Unsupported index dtype: %s

What it means

cuSPARSE matrix indices must be 32- or 64-bit integers; this helper maps NumPy dtypes to cusparseIndexType_t and throws when the dtype's (kind, itemsize) is not in that set.

Source

Thrown at jaxlib/gpu/sparse.cc:47

#include "xla/tsl/python/lib/core/numpy.h"

namespace nb = nanobind;

namespace jax {
namespace JAX_GPU_NAMESPACE {
namespace {

gpusparseIndexType_t DtypeToCuSparseIndexType(const dtype& np_type) {
  static auto* types =
      new absl::flat_hash_map<std::pair<char, int>, gpusparseIndexType_t>({
          {{'u', 2}, GPUSPARSE_INDEX_16U},
          {{'i', 4}, GPUSPARSE_INDEX_32I},
          {{'i', 8}, GPUSPARSE_INDEX_64I},
      });
  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 index dtype: %s", repr.c_str()));
  }
  return it->second;
}

gpuDataType DtypeToCudaDataType(const dtype& np_type) {
  static auto* types =
      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},

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast indices to np.int32 (or np.int64 for large matrices) before the call
2. Check matrix nnz > INT32_MAX and use int64 indices then
3. Upgrade jax/jaxlib — newer versions handle index dtype conversion internally

Example fix

# before
mat.indices.astype(jax.numpy.uint32)
# after
mat.indices.astype(jax.numpy.int32)

idx = np.asarray(indices)
assert idx.dtype in (np.int32, np.int64), f'bad index dtype {idx.dtype}'

def valid_index_dtype(a) -> bool:
    return a.dtype in (np.int32, np.int64)

When it happens

Trigger: Building a sparse matrix descriptor in jaxlib.cudasparse (or jax.experimental.sparse on GPU) with indices dtype of float, bool, uint, or 16-bit int (e.g. indices stored as np.uint32 or bfloat16).

Common situations: Users building BCOO/CSR matrices whose index arrays got upcast or created as unsigned/short ints; older jax versions defaulting indices to unusual dtypes.

Related errors

• Expected metadata dtype to be uint2, got: {meta.dtype}
• Sparse metadata format not implemented for {operand_dtype=}
• Sparse MMA unsupported for f32
• A, X must have same dtypes (were {test_output.dtype}, {dt})
• data types do not match: {data.dtype=} {b.dtype=}

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