jax-ml/jax · error · NotImplementedError

Tiling without swizzle is not supported.

Error message

Tiling without swizzle is not supported.

What it means

Error "Tiling without swizzle is not supported." thrown in jax-ml/jax.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:2150

    transform_avals = [transform_avals[0], *transform_avals[2:]]
  else:
    swizzle = None

  # We verify tiling against swizzling in memory before transposes flip it.
  # Note: we do not check len(tiling) == 2 here because subsequent transforms
  # (e.g. ReshapeTransform) may commute with it to produce 2D tiling.
  tiling_transform = None
  if transforms and isinstance(transforms[0], gpu_core.UntilingTransform):
    tiling_transform = transforms[0]
  elif (
      len(transforms) > 1
      and isinstance(transforms[0], gpu_core.ExtractAliasedRef)
      and isinstance(transforms[1], gpu_core.UntilingTransform)
  ):
    tiling_transform = transforms[1]
  if tiling_transform is not None:
    if swizzle is None:
      raise NotImplementedError("Tiling without swizzle is not supported.")
    bw = dtypes.itemsize_bits(ctx.avals_out[0].dtype)
    expected_minor_tiling = swizzle * 8 // bw
    if tiling_transform.tiling[-1] != expected_minor_tiling:
      raise NotImplementedError(
          "Minor tiling dimension does not fit swizzle: "
          f" expected {expected_minor_tiling}, got {tiling_transform.tiling[-1]}"
      )

  x_smem, _, transforms = _handle_transforms(
      ctx, ctx.avals_in[0], x_ref, transform_avals, transforms,
      allow_peer_refs=True
  )
  del x_ref  # Don't use x_ref anymore. Use x_smem instead!

  is_signed = mgpu_utils.is_signed(dtype)

  if not ctx.avals_out[0].shape:  # The scalar case is simple.
    val = memref_dialect.load(x_smem, [])

View on GitHub (pinned to 1e1c6a8fc0)

When it happens

Trigger: Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:2150 when the library encounters an invalid state.

Common situations: See trigger scenarios.

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