{"record":{"id":"df9a21af9ea73e71","repo":"jax-ml/jax","slug":"cannot-commute-unswizzle-and-indexer-with-aval","errorCode":null,"errorMessage":"Cannot commute unswizzle and indexer with {aval}, which does not have a dtype","messagePattern":"Cannot commute unswizzle and indexer with (.+?), which does not have a dtype","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1239,"sourceCode":"      raise ValueError(\"Can't transpose the swizzled dimension.\")\n    return transpose, self\n\n  def commute_reshape(\n      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform\n  ) -> tuple[state_types.ReshapeTransform, UnswizzleRef]:\n    shape = aval.shape\n    if shape[-1] != self.swizzle_elems(aval.dtype):\n      raise ValueError(\n          f\"Reshape shape {shape} is not divisible by swizzle elements\"\n          f\" {self.swizzle_elems(aval.dtype)}\"\n      )\n    return transform, self\n\n  def commute_ndindexer(\n      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer\n  ) -> tuple[indexing.NDIndexer, UnswizzleRef]:\n    if not hasattr(aval, \"dtype\"):\n      raise ValueError(\n          f\"Cannot commute unswizzle and indexer with {aval}, which does not\"\n          \" have a dtype\"\n      )\n    dtype = aval.dtype\n    swizzle_elems = self.swizzle_elems(dtype)\n    idxs = indexer.indices\n    if not idxs:\n      return indexer, self\n    if not all(isinstance(idx, (slice, indexing.Slice)) for idx in idxs[-2:]):\n      raise NotImplementedError(\n          f\"Non-slice indices are not supported in 2 minormost dims: {idxs}\"\n      )\n    last_idx = idxs[-1]\n    if isinstance(last_idx, indexing.Slice):\n      if last_idx.start != 0 or last_idx.size != swizzle_elems:\n        raise ValueError(\"Swizzled dims cannot be sliced\")\n    else:\n      assert isinstance(last_idx, slice)","sourceCodeStart":1221,"sourceCodeEnd":1257,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1221-L1257","documentation":"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.","triggerScenarios":"Indexing a swizzled ref whose aval lacks .dtype — typically indexing the AbstractRef aval directly instead of its inner ShapedArray during transform commuting.","commonSituations":"Custom Pallas state types or nested refs passed into swizzled layouts; bugs in user code that swizzles non-array state.","solutions":["Only index swizzled refs whose aval is a ShapedArray (has dtype)","Pass the inner array aval, not the ref aval, when building indexers","Avoid swizzling non-array state"],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":"assert hasattr(aval, 'dtype') and aval.dtype is not None, 'aval needs a dtype for unswizzle commuting'","typeGuard":"def has_dtype(aval) -> bool:\\n    return hasattr(aval, 'dtype')","tryCatchPattern":null,"preventionTips":["Only index swizzled refs backed by ShapedArray avals","Pass inner array avals when constructing indexers"],"tags":["jax","pallas","dtype","swizzle","indexing"],"backgroundTag":"missing-dtype-on-aval","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}