jax-ml/jax · error · ValueError

Cannot apply {tiling_transform_str} with {swizzle_str} to me

Error message

Cannot apply {tiling_transform_str} with {swizzle_str} to memref with shape {source.shape}.

What it means

When applying with_transforms, Mosaic checks that the tiling vector (and optional swizzle) evenly divides the source memref's shape. If any dimension of the memref is not a multiple of the corresponding tiling (or is smaller than the tile), the transforms cannot be applied and a ValueError is raised describing the mismatch.

Source

Thrown at jax/experimental/mosaic/gpu/layout_inference.py:2165

  source = ValueSite(op, VariableType.OPERAND, 0)
  dest = ValueSite(op, VariableType.RESULT, 0)
  var = ctx.producer_ref(source)
  smem_transforms = _extract_smem_transforms_from_custom_transform_attrs(
      op.transforms
  )

  if not cs.is_valid_assignment(var, smem_transforms):
    tiling_transform_str = (
        f"tiling {smem_transforms.tiling}"
        if smem_transforms.tiling
        else "empty tiling"
    )
    swizzle_str = (
        f"{smem_transforms.swizzle} swizzle"
        if smem_transforms.swizzle
        else "no swizzle"
    )
    raise ValueError(
        f"Cannot apply {tiling_transform_str} with {swizzle_str} to memref with"
        f" shape {source.shape}."
    )
  assignments: dict[cs.Variable, cs.Constant] = {var: smem_transforms}
  return cs.ConstraintSystem(assignments=assignments), {var: [source, dest]}


def _vector_value_sites_and_constraints_for_async_ops(
    op: mgpu.AsyncLoadOp | mgpu.AsyncStoreOp | mgpu.AsyncPrefetchOp,
) -> tuple[ValueSitesForVariable, list[cs.Constraint]]:
  values_sites: ValueSitesForVariable = dict()
  constraints: list[cs.Constraint] = []

  match op:
    case mgpu.AsyncLoadOp():
      base_operand_index = 3
    case mgpu.AsyncStoreOp():
      base_operand_index = 2

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the smem/tile shape a multiple of the tiling vector and swizzle factor (e.g. pad the last tile dimension up to the swizzle bytes)
2. Derive the tiling vector from the memref shape: tile dims should divide the corresponding memref dims
3. Reduce the swizzle (e.g. from 128B to 64B/32B) or disable swizzle for small tiles

Example fix

# before: ValueError — 64 not divisible by 128
t = mgpu.smem_transforms(tiling=[128, 128], swizzle=128)
out = mgpu.with_transforms(smem_shaped_64x64, t)

# after
t = mgpu.smem_transforms(tiling=[64, 64], swizzle=64)
out = mgpu.with_transforms(smem_shaped_64x64, t)

tile, sw = smem_transforms.tile, smem_transforms.swizzle or 1
assert all(d % t == 0 and d % sw == 0 for d, t in zip(shape, tile)), 'pad memref to multiple of tiling/swizzle'

def transforms_fit(shape, tiling, swizzle=None) -> bool:
    sw = swizzle or 1
    return all(d % t == 0 and d % sw == 0 for d, t in zip(shape, tiling))

When it happens

Trigger: Calling with_transforms (or a custom primitive carrying smem transforms) where e.g. tiling [128, 128] with swizzle 128 is applied to a memref of shape [64, 64], or a dimension not divisible by the tile/swizzle factor.

Common situations: Hard-coded tile/swizzle sizes (common 128B swizzle) combined with small or odd-shaped tail tiles; kernels parameterized by problem size where small N triggers the check; changing swizzle mode without resizing smem buffers.

Related errors

• Unsupported transforms {transforms}
• Swizzle {self.swizzle} requires the trailing dimension to be
• {tiling=} and {grid=} must have same length.
• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is

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