{"record":{"id":"e2462e638fea3749","repo":"jax-ml/jax","slug":"unsupported-dtype-dtype-e2462e","errorCode":null,"errorMessage":"Unsupported dtype: {dtype}","messagePattern":"Unsupported dtype: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_core.py","lineNumber":212,"sourceCode":"    return [\n        tpu_core.MemorySpace.VMEM,\n        tpu_core.MemorySpace.VMEM_SHARED,\n        tpu_core.MemorySpace.SMEM,\n        tpu_core.MemorySpace.SEMAPHORE,\n    ]\n\n  @contextlib.contextmanager\n  def tracing_context(self):\n    yield\n\n\ndef supported_shapes(dtype: jax.typing.DTypeLike) -> Sequence[tuple[int, ...]]:\n  \"\"\"Returns all supported array shapes for the given dtype on SparseCore.\"\"\"\n  sc_info = get_sparse_core_info()\n  num_lanes = sc_info.num_lanes\n  itemsize = jnp.dtype(dtype).itemsize\n  if itemsize > 4:\n    raise ValueError(f\"Unsupported dtype: {dtype}\")\n  packing_factor = 4 // itemsize\n  if packing_factor == 1:\n    return [(num_lanes,)]\n  return [(num_lanes * packing_factor,), (packing_factor, num_lanes)]\n\n\n@tree_util.register_dataclass\n@dataclasses.dataclass(frozen=True)\nclass Indices:\n  \"\"\"Indices for a gather or a scatter on SparseCore.\n\n  Attributes:\n    values: The values of the indices. Can be an array or a ref.\n    ignored_value: If not None, the indices with this value will be ignored.\n  \"\"\"\n\n  values: Any\n  ignored_value: int | None = jax.tree.static(default=None)","sourceCodeStart":194,"sourceCodeEnd":230,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_core.py#L194-L230","documentation":"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.","triggerScenarios":"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).","commonSituations":"Defaulting kernel buffers to int64 indices on TPU (where jax_enable_x64 matters); porting a CPU/GPU kernel using double precision.","solutions":["Switch operands to 32-bit or smaller dtypes (float32, int32, bfloat16, etc.)","If indices are int64 because of x64 mode, disable jax_enable_x64 or cast to int32 before the kernel","Check jnp.dtype(x).itemsize <= 4 for every SC operand"],"exampleFix":"# before\nkernel(pltpu.HBM(x.astype(jnp.float64), ...))\n# after\nkernel(pltpu.HBM(x.astype(jnp.float32), ...))","handlingStrategy":"validation","validationCode":"assert jnp.dtype(dtype).itemsize <= 4, f'Unsupported dtype {dtype}'","typeGuard":"def sc_supported_dtype(dt) -> bool:\n    import numpy as np\n    return np.dtype(dt).itemsize <= 4","tryCatchPattern":"null","preventionTips":["Cast all SC operands to 32-bit or smaller","Keep jax_enable_x64=False for SC kernels"],"tags":["jax","pallas","tpu","sparsecore","dtype"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}