{"record":{"id":"ce0711037aca5dbc","repo":"jax-ml/jax","slug":"subclasses-should-override-this-method","errorCode":null,"errorMessage":"Subclasses should override this method","messagePattern":"Subclasses should override this method","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":120,"sourceCode":"\n@dataclasses.dataclass(frozen=True)\nclass _Partitioned(CopyPartition):\n  axis: int\n\n\n@dataclasses.dataclass(frozen=True)\nclass _Replicated(CopyPartition):\n  pass\n\n\nCopyPartition.PARTITIONED = _Partitioned\nCopyPartition.REPLICATED = _Replicated()\n\n\n@dataclasses.dataclass(frozen=True)\nclass MemRefTransform:\n  def apply(self, ref: ir.Value) -> ir.Value:\n    raise NotImplementedError(\"Subclasses should override this method\")\n\n  def transform_index(self, idx: Sequence[ir.Value]) -> tuple[ir.Value, ...]:\n    raise NotImplementedError(\"Subclasses should override this method\")\n\n  def transform_shape(self, shape: Sequence[int]) -> tuple[int, ...]:\n    raise NotImplementedError(\"Subclasses should override this method\")\n\n  def transform_gmem_shape(self, shape: Sequence[int]) -> tuple[int, ...]:\n    \"\"\"Applies the shape transformation to the given GMEM shape.\n\n    This function is intended to mirror the behavior of the `apply` method on\n    GMEM shapes.\n    \"\"\"\n    raise NotImplementedError(\"Subclasses should override this method\")\n\n  def transform_strides(self, strides: Sequence[int]) -> tuple[int, ...]:\n    raise NotImplementedError(\"Subclasses should override this method\")\n","sourceCodeStart":102,"sourceCodeEnd":138,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L102-L138","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Override apply(self, ref: ir.Value) -> ir.Value in your MemRefTransform subclass","Use an existing concrete subclass instead of the base class","Mark custom base subclasses with abc.ABC/abstractmethod to fail at instantiation rather than call time"],"exampleFix":"# before\nclass MyTransform(MemRefTransform):\n  pass  # apply not overridden -> NotImplementedError when called\n\n# after\nclass MyTransform(MemRefTransform):\n  def apply(self, ref: ir.Value) -> ir.Value:\n    return ...  # transform the memref value","handlingStrategy":"validation","validationCode":"def transform_complete(t) -> bool:\n  required = ['apply', 'transform_index', 'transform_shape', 'transform_gmem_shape', 'transform_strides']\n  return all(\n      getattr(type(t), m, MemRefTransform.__dict__[m]) is not MemRefTransform.__dict__[m]\n      for m in required\n  )\nassert transform_complete(my_transform)","typeGuard":"def has_apply_override(t) -> bool:\n  return type(t).apply is not MemRefTransform.apply","tryCatchPattern":"try:\n  ctx.init_tma_desc(...)\nexcept NotImplementedError as e:\n  if 'Subclasses should override' in str(e):\n    raise TypeError(f'{type(t).__name__} is missing a MemRefTransform override') from e\n  raise","preventionTips":["Subclass MemRefTransform and override all methods used by your path (apply at minimum)","Prefer built-in transforms (transpose/slice) over custom ones","Add unit tests that run TMA lowering with your custom transform"],"tags":["mosaic-gpu","not-implemented","abstract-method","memref-transform","subclassing"],"backgroundTag":"abstract-method-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}