jax-ml/jax · error · ValueError

Found multiple SwizzleTransformAttr

Error message

Found multiple SwizzleTransformAttr

What it means

A memref's in/out_transforms attribute contained more than one SwizzleTransformAttr; Mosaic allows at most one swizzle transform per memory operation.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:973

    _: LoweringContext, op: mgpu.VectorConcatOp
) -> Sequence[ir.Value]:
  in_layouts = inference_utils.in_layouts(op)
  out_layout, = inference_utils.out_layouts(op)
  operands_fa = [
      _fragmented_array_from_ir(opr, l)
      for opr, l in zip(op.operands, in_layouts, strict=True)
  ]
  out = fa.concatenate(operands_fa, axis=op.dimension.value)
  assert out.layout == layouts_lib.from_layout_attr(out_layout)
  return [fragmented_array_to_ir(out, op.result.type)]


def swizzle_from_transforms_attr(attr: ir.ArrayAttr) -> mgpu.SwizzlingMode:
  swizzle = None
  for transform in attr:
    if isinstance(transform, mgpu.SwizzleTransformAttr):
      if swizzle is not None:
        raise ValueError("Found multiple SwizzleTransformAttr")
      swizzle = mgpu.SwizzlingMode(mgpu.SwizzleTransformAttr(transform).swizzle)
  return swizzle or mgpu.SwizzlingMode.kNoSwizzle


def memref_transforms_from_transforms_attr(
    attr: ir.ArrayAttr,
) -> tuple[lc.MemRefTransform, ...]:
  gmem_transforms: list[lc.MemRefTransform] = []
  for transform in attr:
    if isinstance(transform, mgpu.TileTransformAttr):
      tile_transform = lc.TileTransform(tuple(transform.tiling))
      gmem_transforms.append(tile_transform)
    elif not isinstance(transform, mgpu.SwizzleTransformAttr):
      raise NotImplementedError(f"Unsupported transform: {transform}")
  return tuple(gmem_transforms)


def tile_offset(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure exactly zero or one SwizzleTransformAttr in the transforms array
2. Build the attribute via helper code that deduplicates swizzle entries

Example fix

// before
transforms = [swizzle_128b, swizzle_128b, tile(8, 128)]
// after
transforms = [swizzle_128b, tile(8, 128)]

swizzles = [t for t in transforms if isinstance(t, mgpu.SwizzleTransformAttr)]
assert len(swizzles) <= 1, 'at most one swizzle transform allowed'

When it happens

Trigger: Attaching two SwizzleTransformAttr entries to the same load/store's transforms array attribute.

Common situations: Manually building transforms arrays or double-appending swizzle when constructing mgpu ops programmatically.

Related errors

• Unsupported transform: {transform}
• Swizzle={b_swizzle} is too big for MMA with M=64. Try loweri
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}
• Unsupported memory space: {orig_ref_ty.memory_space}

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