{"record":{"id":"ccfea931474f938b","repo":"jax-ml/jax","slug":"transposed-memrefs-are-not-supported-in-expandshap","errorCode":null,"errorMessage":"Transposed memrefs are not supported in ExpandShapeOp.","messagePattern":"Transposed memrefs are not supported in ExpandShapeOp\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2016,"sourceCode":"      ir.AffineDimExpr(e).position for e in op.permutation.value.results\n  )\n  inv_permutation = tuple(permutation.index(i) for i in range(len(permutation)))\n\n  constraints = [\n      cs.Equals(cs.Transpose(source_var, permutation=permutation), dest_var),\n      cs.Equals(source_var, cs.Transpose(dest_var, permutation=inv_permutation)),\n  ]\n  system = cs.ConstraintSystem(constraints=constraints)\n  return system, {source_var: [source], dest_var: [dest]}\n\n\n@_add_constraint_system_derivation_rule(memref.ExpandShapeOp)\ndef _memref_expand_shape_op_equation_system(\n    ctx: DerivationContext,\n    op: memref.ExpandShapeOp,\n) -> ConstraintSystemDerivationRuleResult:\n  if utils.is_memref_transposed(ir.MemRefType(op.src.type)):\n    raise NotImplementedError(\n        \"Transposed memrefs are not supported in ExpandShapeOp.\"\n    )\n\n  source = ValueSite(op, VariableType.OPERAND, 0)\n  source_var = ctx.producer_ref(source)\n  dest = ValueSite(op, VariableType.RESULT, 0)\n  dest_var = cs.Variable(dest)\n\n  reverse_tiling_multiple = []\n  for dim, idx in zip(\n      reversed(op.static_output_shape), reversed(op.reassociation)\n  ):\n    # pyrefly: ignore[bad-argument-type]\n    if ir.ShapedType.is_dynamic_size(dim) or len(idx) > 1:\n      # For simplicity, we only support tiling non-expanded static dimensions.\n      # These limitations could be lifted later if needed.\n      break\n    reverse_tiling_multiple.append(dim)","sourceCodeStart":1998,"sourceCodeEnd":2034,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L1998-L2034","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\nsv = memref.transpose(src)          # non-contiguous\nout = memref.expand_shape(sv, reassoc=[[0,1],[2]])\n\n# after\nexp = memref.expand_shape(src, reassoc=[[0,1],[2]])  # contiguous source\nout = memref.transpose(exp)","handlingStrategy":"validation","validationCode":"from jax.experimental.mosaic.gpu import layout_inference_utils as utils\nif utils.is_memref_transposed(ir.MemRefType(src.type)):\n    # copy or reorder before expand_shape\n    src = make_contiguous(src)","typeGuard":"def is_expand_safe(memref_value) -> bool:\n    m = ir.MemRefType(memref_value.type)\n    return m.get_strides_and_offset()[0] == utils.get_contiguous_strides(m.shape)","tryCatchPattern":null,"preventionTips":["Keep expand_shape/collapse_shape immediately adjacent to allocations, before any transpose or strided view","Centralize reshapes in helpers that assert contiguous strides"],"tags":["jax","mosaic-gpu","memref","expand-shape","transpose","layout-inference"],"backgroundTag":"unsupported-operation-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}