jax-ml/jax · error · NotImplementedError

Expanding tiled dimensions is not supported.

Error message

Expanding tiled dimensions is not supported.

What it means

memref.expand_shape lowering cannot split a tiled dimension into multiple dimensions: if any of the last num_tiling_dims reassociation groups has more than one member, expansion of tiled dims is rejected.

Source

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

  in_transforms = inference_utils.in_transforms(op)[0]
  unwrapped_in_ref = unwrap_transformed_memref(op.src, in_transforms)
  in_transformed_ty = ir.MemRefType(unwrapped_in_ref.type)

  out_transforms = inference_utils.out_transforms(op)[0]
  out_transformed_ty = transform_type(op.result.type, out_transforms)

  reassociation = cast(list[ir.ArrayAttr], list(op.reassociation))
  num_tiling_dims = len(in_transformed_ty.shape) - len(op.src.type.shape)

  # We don't currently allow expanding tiled dimensions. So to compute the
  # reassociation on the lowered types, we just need to backfill the original
  # one with the number of missing dimensions.
  if num_tiling_dims > 0 and any(
      len(x) > 1 for x in reassociation[-num_tiling_dims:]
  ):
    # If we ever remove this restriction, we will need to ensure this is
    # compatible with `transform_type`.
    raise NotImplementedError("Expanding tiled dimensions is not supported.")

  start_index = len(op.static_output_shape)
  for i in range(start_index, start_index + num_tiling_dims):
    reassociation.append([i])  # pyrefly: ignore[bad-argument-type]

  new_expand_shape_op = memref.ExpandShapeOp(
      out_transformed_ty,
      unwrapped_in_ref,
      reassociation,
      output_shape=op.output_shape,
      static_output_shape=out_transformed_ty.shape,
  )

  wrapped_ref = wrap_transformed_memref(
      new_expand_shape_op.result, op.result.type, out_transforms
  )
  return [wrapped_ref]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Expand only untiled dimensions; keep tiled dims intact (reassociation entry [i] alone)
2. Retile with smaller tiles before expanding so expansion targets untiled dims
3. Do the expansion before applying the tile transform

Example fix

// before
# tiled last dim of size 8
out = t.memref.expand_shape(tiled, reassociation=[[0],[1,2]])  # splits tiled dim
// after
out = t.memref.expand_shape(tiled, reassociation=[[0,1],[2]])  # split untiled dims only

assert all(len(g) == 1 for g in reassociation[-num_tiling_dims:]), 'cannot expand tiled dims'

When it happens

Trigger: memref.expand_shape on a tiled memref where a reassociation group covering a tiled dimension expands it into 2+ output dims (len(group) > 1).

Common situations: Reshaping a tiled smem tensor to expose inner axes (e.g. splitting the last tiled dim into two) for element-wise access; not supported due to transform_type incompatibility noted in the source comment.

Related errors

• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Folding tiled dimensions into untiled dimensions is not supp

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