jax-ml/jax · error · NotImplementedError

Transposed memrefs are not supported in ExpandShapeOp.

Error message

Transposed memrefs are not supported in ExpandShapeOp.

What it means

Mosaic GPU's layout inference handles memref.ExpandShapeOp only for non-transposed source memrefs. If the source memref's strides do not match the contiguous row-major stride pattern (i.e. is_memref_transposed reports a transpose), the constraint derivation bails out with NotImplementedError because the expand-shape semantics under a transpose are not modeled.

Source

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

      ir.AffineDimExpr(e).position for e in op.permutation.value.results
  )
  inv_permutation = tuple(permutation.index(i) for i in range(len(permutation)))

  constraints = [
      cs.Equals(cs.Transpose(source_var, permutation=permutation), dest_var),
      cs.Equals(source_var, cs.Transpose(dest_var, permutation=inv_permutation)),
  ]
  system = cs.ConstraintSystem(constraints=constraints)
  return system, {source_var: [source], dest_var: [dest]}


@_add_constraint_system_derivation_rule(memref.ExpandShapeOp)
def _memref_expand_shape_op_equation_system(
    ctx: DerivationContext,
    op: memref.ExpandShapeOp,
) -> ConstraintSystemDerivationRuleResult:
  if utils.is_memref_transposed(ir.MemRefType(op.src.type)):
    raise NotImplementedError(
        "Transposed memrefs are not supported in ExpandShapeOp."
    )

  source = ValueSite(op, VariableType.OPERAND, 0)
  source_var = ctx.producer_ref(source)
  dest = ValueSite(op, VariableType.RESULT, 0)
  dest_var = cs.Variable(dest)

  reverse_tiling_multiple = []
  for dim, idx in zip(
      reversed(op.static_output_shape), reversed(op.reassociation)
  ):
    # pyrefly: ignore[bad-argument-type]
    if ir.ShapedType.is_dynamic_size(dim) or len(idx) > 1:
      # For simplicity, we only support tiling non-expanded static dimensions.
      # These limitations could be lifted later if needed.
      break
    reverse_tiling_multiple.append(dim)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Move the expand_shape before the transpose so it operates on a contiguous memref, then transpose the expanded result
Copy the transposed memref into a fresh contiguous buffer (e.g. via a load/store round-trip through registers) before expanding
Re-express the reshape as slice + concat or explicit copy loops that keep contiguity

Example fix

# before
sv = memref.transpose(src)          # non-contiguous
out = memref.expand_shape(sv, reassoc=[[0,1],[2]])

# after
exp = memref.expand_shape(src, reassoc=[[0,1],[2]])  # contiguous source
out = memref.transpose(exp)

Defensive patterns

Strategy: validation

Validate before calling

from jax.experimental.mosaic.gpu import layout_inference_utils as utils
if utils.is_memref_transposed(ir.MemRefType(src.type)):
    # copy or reorder before expand_shape
    src = make_contiguous(src)

Type guard

def is_expand_safe(memref_value) -> bool:
    m = ir.MemRefType(memref_value.type)
    return m.get_strides_and_offset()[0] == utils.get_contiguous_strides(m.shape)

Prevention

Keep expand_shape/collapse_shape immediately adjacent to allocations, before any transpose or strided view
Centralize reshapes in helpers that assert contiguous strides

When it happens

Trigger: Applying memref.expand_shape (or an equivalent reshape that expands a dimension) to a memref whose layout was produced by memref.transpose, a strided view, or a swizzled shared-memory layout, inside a Mosaic GPU kernel.

Common situations: Kernels that transpose tiles in shared memory and then reshape; reshape after tiling/slicing ops that leave non-contiguous strides. Common when porting Triton-style kernels where transpose+reshape is idiomatic.

Related errors

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