jax-ml/jax · error · ValueError

Unsupported dtype: {dtype}

Error message

Unsupported dtype: {dtype}

What it means

supported_shapes() enumerates the array shapes SparseCore kernels can operate on given a dtype; the SC vector lanes are 4 bytes, so dtypes with itemsize > 4 (e.g. float64, int64) cannot be packed and are rejected.

Source

Thrown at jax/_src/pallas/mosaic/sc_core.py:212

    return [
        tpu_core.MemorySpace.VMEM,
        tpu_core.MemorySpace.VMEM_SHARED,
        tpu_core.MemorySpace.SMEM,
        tpu_core.MemorySpace.SEMAPHORE,
    ]

  @contextlib.contextmanager
  def tracing_context(self):
    yield


def supported_shapes(dtype: jax.typing.DTypeLike) -> Sequence[tuple[int, ...]]:
  """Returns all supported array shapes for the given dtype on SparseCore."""
  sc_info = get_sparse_core_info()
  num_lanes = sc_info.num_lanes
  itemsize = jnp.dtype(dtype).itemsize
  if itemsize > 4:
    raise ValueError(f"Unsupported dtype: {dtype}")
  packing_factor = 4 // itemsize
  if packing_factor == 1:
    return [(num_lanes,)]
  return [(num_lanes * packing_factor,), (packing_factor, num_lanes)]


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class Indices:
  """Indices for a gather or a scatter on SparseCore.

  Attributes:
    values: The values of the indices. Can be an array or a ref.
    ignored_value: If not None, the indices with this value will be ignored.
  """

  values: Any
  ignored_value: int | None = jax.tree.static(default=None)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Switch operands to 32-bit or smaller dtypes (float32, int32, bfloat16, etc.)
2. If indices are int64 because of x64 mode, disable jax_enable_x64 or cast to int32 before the kernel
3. Check jnp.dtype(x).itemsize <= 4 for every SC operand

Example fix

# before
kernel(pltpu.HBM(x.astype(jnp.float64), ...))
# after
kernel(pltpu.HBM(x.astype(jnp.float32), ...))

assert jnp.dtype(dtype).itemsize <= 4, f'Unsupported dtype {dtype}'

def sc_supported_dtype(dt) -> bool:
    import numpy as np
    return np.dtype(dt).itemsize <= 4

null

When it happens

Trigger: Passing a 64-bit dtype (jnp.float64, jnp.int64, jnp.uint64) to supported_shapes, or using such a dtype in a SparseCore Pallas kernel operand (surfaced via _check_aval_is_supported).

Common situations: Defaulting kernel buffers to int64 indices on TPU (where jax_enable_x64 matters); porting a CPU/GPU kernel using double precision.

Related errors

• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• Unsigned integer dtype {aval.dtype} is not supported for dot

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