jax-ml/jax · error · NotImplementedError

Subclasses should implement this method.

Error message

Subclasses should implement this method.

What it means

Base Sharding.device_set() raises NotImplementedError: the base jax.sharding.Sharding class does not implement it; only concrete subclasses (NamedSharding, SingleDeviceSharding, GSPMDSharding, ...) do. Hitting it means you instantiated the abstract base or an incomplete subclass.

Source

Thrown at jax/_src/sharding.py:106

            s1._internal_device_list == s2._internal_device_list)
  else:
    return hlo_s_eq and mem_eq


@use_cpp_class(xc.Sharding)
class Sharding:
  """Describes how a :class:`jax.Array` is laid out across devices.
  """

  # Abstract methods below that subclasses should implement.
  @property
  def device_set(self) -> set[Device]:
    """The set of devices that this :class:`Sharding` spans.

    In multi-controller JAX, the set of devices is global, i.e., includes
    non-addressable devices from other processes.
    """
    raise NotImplementedError('Subclasses should implement this method.')

  @property
  def is_fully_replicated(self) -> bool:
    """Is this sharding fully replicated?

    A sharding is fully replicated if each device has a complete copy of the
    entire data.
    """
    raise NotImplementedError('Subclasses should implement this method.')

  @property
  def is_fully_addressable(self) -> bool:
    """Is this sharding fully addressable?

    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.
    """

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a concrete sharding (NamedSharding, SingleDeviceSharding, PositionalSharding, GSPMDSharding) instead of the base class
2. If subclassing, implement device_set (and the other abstract members) in your subclass

Example fix

# before
sh = jax.sharding.Sharding()
sh.device_set()

# after
from jax.sharding import NamedSharding, PartitionSpec as P
sh = NamedSharding(mesh, P('data'))
sh.device_set()

from jax.sharding import Sharding, NamedSharding
assert not type(s) is Sharding, 'base Sharding is abstract'

def is_concrete_sharding(s) -> bool:
    from jax.sharding import Sharding
    return isinstance(s, Sharding) and type(s) is not Sharding and getattr(type(s).device_set, '__isabstractmethod__', False) is False

When it happens

Trigger: Calling Sharding().device_set(), or subclassing Sharding without overriding device_set, then calling it (often indirectly via addressable_devices).

Common situations: Custom Sharding subclass missing required methods; accidentally returning base-class instances from factory code.

Related errors

• 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
• Mapped away dimension of inputs passed to vmap should be sha

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