jax-ml/jax · error · NotImplementedError

Subclasses should override this method

Error message

Subclasses should override this method

What it means

MemRefTransform is an abstract base (frozen dataclass) for transforms applied to memref values during TMA/async-copy lowering. Its apply(ref) method raises NotImplementedError; concrete subclasses (e.g. TransposeMemRefTransform, SliceMemRefTransform) must override it. Hitting this means a transform instance without an apply override was used by init_tma_desc/_prepare_tma/async_copy.

Source

Thrown at jax/experimental/mosaic/gpu/launch_context.py:120

@dataclasses.dataclass(frozen=True)
class _Partitioned(CopyPartition):
  axis: int


@dataclasses.dataclass(frozen=True)
class _Replicated(CopyPartition):
  pass


CopyPartition.PARTITIONED = _Partitioned
CopyPartition.REPLICATED = _Replicated()


@dataclasses.dataclass(frozen=True)
class MemRefTransform:
  def apply(self, ref: ir.Value) -> ir.Value:
    raise NotImplementedError("Subclasses should override this method")

  def transform_index(self, idx: Sequence[ir.Value]) -> tuple[ir.Value, ...]:
    raise NotImplementedError("Subclasses should override this method")

  def transform_shape(self, shape: Sequence[int]) -> tuple[int, ...]:
    raise NotImplementedError("Subclasses should override this method")

  def transform_gmem_shape(self, shape: Sequence[int]) -> tuple[int, ...]:
    """Applies the shape transformation to the given GMEM shape.

    This function is intended to mirror the behavior of the `apply` method on
    GMEM shapes.
    """
    raise NotImplementedError("Subclasses should override this method")

  def transform_strides(self, strides: Sequence[int]) -> tuple[int, ...]:
    raise NotImplementedError("Subclasses should override this method")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Override apply(self, ref: ir.Value) -> ir.Value in your MemRefTransform subclass
2. Use an existing concrete subclass instead of the base class
3. Mark custom base subclasses with abc.ABC/abstractmethod to fail at instantiation rather than call time

Example fix

# before
class MyTransform(MemRefTransform):
  pass  # apply not overridden -> NotImplementedError when called

# after
class MyTransform(MemRefTransform):
  def apply(self, ref: ir.Value) -> ir.Value:
    return ...  # transform the memref value

def transform_complete(t) -> bool:
  required = ['apply', 'transform_index', 'transform_shape', 'transform_gmem_shape', 'transform_strides']
  return all(
      getattr(type(t), m, MemRefTransform.__dict__[m]) is not MemRefTransform.__dict__[m]
      for m in required
  )
assert transform_complete(my_transform)

def has_apply_override(t) -> bool:
  return type(t).apply is not MemRefTransform.apply

try:
  ctx.init_tma_desc(...)
except NotImplementedError as e:
  if 'Subclasses should override' in str(e):
    raise TypeError(f'{type(t).__name__} is missing a MemRefTransform override') from e
  raise

When it happens

Trigger: Subclassing MemRefTransform (or instantiating it directly) without overriding apply, then passing the transform to launch-context APIs like init_tma_desc or async_copy which call transform.apply(ref).

Common situations: Writing a custom MemRefTransform for a new copy pattern and forgetting one required method; refactors that renamed the override or changed its signature; instantiating the base class directly instead of a concrete subclass.

Related errors

• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Sparse metadata format not implemented for {operand_dtype=}

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