jax-ml/jax · error · ValueError

Cannot commute unswizzle and indexer with {aval}, which does

Error message

Cannot commute unswizzle and indexer with {aval}, which does not have a dtype

What it means

UnswizzleRef.commute_ndindexer requires the aval it commutes with to expose a dtype (to compute swizzle_elems). If the aval has no dtype attribute (e.g. a token or abstract ref itself rather than its inner array), the transform cannot determine the swizzle granularity and raises ValueError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1239

      raise ValueError("Can't transpose the swizzled dimension.")
    return transpose, self

  def commute_reshape(
      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform
  ) -> tuple[state_types.ReshapeTransform, UnswizzleRef]:
    shape = aval.shape
    if shape[-1] != self.swizzle_elems(aval.dtype):
      raise ValueError(
          f"Reshape shape {shape} is not divisible by swizzle elements"
          f" {self.swizzle_elems(aval.dtype)}"
      )
    return transform, self

  def commute_ndindexer(
      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, UnswizzleRef]:
    if not hasattr(aval, "dtype"):
      raise ValueError(
          f"Cannot commute unswizzle and indexer with {aval}, which does not"
          " have a dtype"
      )
    dtype = aval.dtype
    swizzle_elems = self.swizzle_elems(dtype)
    idxs = indexer.indices
    if not idxs:
      return indexer, self
    if not all(isinstance(idx, (slice, indexing.Slice)) for idx in idxs[-2:]):
      raise NotImplementedError(
          f"Non-slice indices are not supported in 2 minormost dims: {idxs}"
      )
    last_idx = idxs[-1]
    if isinstance(last_idx, indexing.Slice):
      if last_idx.start != 0 or last_idx.size != swizzle_elems:
        raise ValueError("Swizzled dims cannot be sliced")
    else:
      assert isinstance(last_idx, slice)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Only index swizzled refs whose aval is a ShapedArray (has dtype)
2. Pass the inner array aval, not the ref aval, when building indexers
3. Avoid swizzling non-array state

assert hasattr(aval, 'dtype') and aval.dtype is not None, 'aval needs a dtype for unswizzle commuting'

def has_dtype(aval) -> bool:\n    return hasattr(aval, 'dtype')

When it happens

Trigger: Indexing a swizzled ref whose aval lacks .dtype — typically indexing the AbstractRef aval directly instead of its inner ShapedArray during transform commuting.

Common situations: Custom Pallas state types or nested refs passed into swizzled layouts; bugs in user code that swizzles non-array state.

Related errors

• Non-slice indices are not supported in 2 minormost dims: {id
• Indexer must have integer or boolean type, got indexer with
• Indexer must have integer or boolean type, got indexer with
• Indexer must have integer or boolean type, got indexer with
• np.delete(arr, obj): got obj.dtype={obj_array.dtype}; must b

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