jax-ml/jax · error · NotImplementedError

Subclasses should implement this method

Error message

Subclasses should implement this method

What it means

Generic NotImplementedError sentinel raised by the abstract Sharding base class (here the with_memory_kind property and similar abstract members) when a subclass fails to override a required method. Indicates a sharding implementation bug, not user input.

Solutions

  1. Use a built-in sharding (NamedSharding/PositionalSharding support with_memory_kind)
  2. Implement with_memory_kind in your subclass, returning a copy configured for the kind

Example fix

# before
sh.with_memory_kind('pinned')  # sh is custom, no override

# after
class MySh(jax.sharding.Sharding):
  def with_memory_kind(self, kind):
    return replace(self, _memory_kind=kind)
Defensive patterns

Strategy: type-guard

Validate before calling

assert type(s).with_memory_kind is not jax.sharding.Sharding.with_memory_kind, 'with_memory_kind not implemented'

Type guard

def has_with_memory_kind_impl(s) -> bool:
    return type(s).with_memory_kind is not jax.sharding.Sharding.with_memory_kind

Prevention

When it happens

Trigger: Calling sharding.with_memory_kind('pinned') on the base class or a custom subclass lacking the override; common when moving arrays to host memory via device_put with a sharding.

Common situations: Custom sharding types used with jax.device_put(..., memory_kind='pinned') flows.

Understand the failure class

Background: "NotImplementedError: Subclasses should override this method" / "must be implemented" — abstract method errors explained — this error's family across 40 libraries.

Related errors

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

Appendix: source

Thrown at jax/_src/sharding.py:139

    A sharding is fully addressable if the current process can address all of
    the devices named in the :class:`Sharding`. ``is_fully_addressable`` is
    equivalent to "is_local" in multi-process JAX.
    """
    raise NotImplementedError('Subclasses should implement this method.')

  @property
  def num_devices(self) -> int:
    """Number of devices that the sharding contains."""
    raise NotImplementedError('Subclasses should implement this method.')

  @property
  def memory_kind(self) -> str | None:
    """Returns the memory kind of the sharding."""
    raise NotImplementedError('Subclasses should implement this method.')

  def with_memory_kind(self, kind: str) -> Sharding:
    """Returns a new Sharding instance with the specified memory kind."""
    raise NotImplementedError('Subclasses should implement this method')

  @property
  def _device_assignment(self) -> XLADeviceAssignment:
    raise NotImplementedError('Subclasses should implement this method.')

  @property
  def _internal_device_list(self) -> xc.DeviceList:
    raise NotImplementedError('Subclasses should implement this method.')

  def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
    raise NotImplementedError('Subclasses should implement this method.')

  def _to_sdy_sharding(self, num_dimensions: int,
                       modify_wrt_axis_types: bool = False):
    raise NotImplementedError('Subclasses should implement this method.')

  #############################################################################
  # Default implementations below that all subclasses will inherit.

View on GitHub (pinned to 1e1c6a8fc0)