jax-ml/jax · error · NotImplementedError

Subclasses should implement this method

Error message

Subclasses should implement this method

What it means

with_memory_kind(kind) on the base jax.sharding.Sharding is abstract: it should return a new sharding instance targeting the requested memory kind, but only concrete subclasses implement it.

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)

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)

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

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

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.

Related errors

• sharding with memory_kind is not allowed. Please use `jax.de
• unreduced rule for {name} is not implemented. Please file an
• Explicit sharding is not currently supported in the pallas-t
• pallas_call with a mesh does not support batching
• Subclasses should implement this method.

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