jax-ml/jax · error · NotImplementedError

memref.StoreOp does not support transforms: {op}

Error message

memref.StoreOp does not support transforms: {op}

What it means

memref.store, like load, is a scalar pass-through in Mosaic lowering and rejects any non-empty transform annotation on the target memref.

Source

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

      indices=op.indices,
      nontemporal=op.nontemporal,
  )
  return [new_load_op.result]


@_register_lowering(memref.StoreOp, support_warp_semantics=True)
def _memref_store_op_lowering_rule(
    ctx: LoweringContext, op: memref.StoreOp
) -> Sequence[ir.Value]:
  """Lowering rule for memref.StoreOp.

  Stores are never transformed so this rule is mostly just a pass-through.
  """
  del ctx

  in_transforms = inference_utils.in_transforms(op)[0]
  if in_transforms:
    raise NotImplementedError(f"memref.StoreOp does not support transforms: {op}")

  memref.StoreOp(
      value=op.value,
      memref=unwrap_transformed_memref(op.memref, in_transforms),
      indices=op.indices,
      nontemporal=op.nontemporal,
  )
  return []


@_register_lowering(mgpu.TmemAllocOp)
def _tmem_alloc_op_lowering_rule(
    ctx: LoweringContext, op: mgpu.TmemAllocOp
) -> Sequence[ir.Value]:
  """Lowering rule for mgpu.TmemAllocOp."""
  ctx.check_collective(op)

  output_shape = ir.MemRefType(op.result.type).shape

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Unwrap the transform and store to the underlying base memref
2. Use store_tensor for transformed buffers
3. Keep memrefs used for scalar stores untransformed

Example fix

// before
t.memref.store(val, tiled_ref, idx)
// after
base = unwrap_transformed_memref(tiled_ref, transforms)
t.memref.store(val, base, idx)

assert not inference_utils.in_transforms(op)[0], 'store does not support transforms'

def storable_untransformed(op) -> bool:
    return not inference_utils.in_transforms(op)[0]

When it happens

Trigger: Emitting memref.store where the stored-to memref has non-empty in_transforms (tiled/swizzled layout still attached).

Common situations: Writing scalar results into a tiled shared-memory buffer directly; mixing element-wise stores with tensor-mem transforms in a fused kernel.

Related errors

• Unsupported transform: {type(transform)}
• Non-indexing transforms on GMEM refs are not implemented.
• Not all transforms could be handled. Remaining transforms: {
• memref.LoadOp does not support transforms: {op}
• Unsupported dtype: {ref.dtype}

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