{"record":{"id":"688402849adeb876","repo":"jax-ml/jax","slug":"subclasses-should-implement-this-method-688402","errorCode":null,"errorMessage":"Subclasses should implement this method","messagePattern":"Subclasses should implement this method","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/sharding.py","lineNumber":139,"sourceCode":"    A sharding is fully addressable if the current process can address all of\n    the devices named in the :class:`Sharding`. ``is_fully_addressable`` is\n    equivalent to \"is_local\" in multi-process JAX.\n    \"\"\"\n    raise NotImplementedError('Subclasses should implement this method.')\n\n  @property\n  def num_devices(self) -> int:\n    \"\"\"Number of devices that the sharding contains.\"\"\"\n    raise NotImplementedError('Subclasses should implement this method.')\n\n  @property\n  def memory_kind(self) -> str | None:\n    \"\"\"Returns the memory kind of the sharding.\"\"\"\n    raise NotImplementedError('Subclasses should implement this method.')\n\n  def with_memory_kind(self, kind: str) -> Sharding:\n    \"\"\"Returns a new Sharding instance with the specified memory kind.\"\"\"\n    raise NotImplementedError('Subclasses should implement this method')\n\n  @property\n  def _device_assignment(self) -> XLADeviceAssignment:\n    raise NotImplementedError('Subclasses should implement this method.')\n\n  @property\n  def _internal_device_list(self) -> xc.DeviceList:\n    raise NotImplementedError('Subclasses should implement this method.')\n\n  def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:\n    raise NotImplementedError('Subclasses should implement this method.')\n\n  def _to_sdy_sharding(self, num_dimensions: int,\n                       modify_wrt_axis_types: bool = False):\n    raise NotImplementedError('Subclasses should implement this method.')\n\n  #############################################################################\n  # Default implementations below that all subclasses will inherit.","sourceCodeStart":121,"sourceCodeEnd":157,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/sharding.py#L121-L157","documentation":"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.","triggerScenarios":"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.","commonSituations":"Custom sharding types used with jax.device_put(..., memory_kind='pinned') flows.","solutions":["Use a built-in sharding (NamedSharding/PositionalSharding support with_memory_kind)","Implement with_memory_kind in your subclass, returning a copy configured for the kind"],"exampleFix":"# before\nsh.with_memory_kind('pinned')  # sh is custom, no override\n\n# after\nclass MySh(jax.sharding.Sharding):\n  def with_memory_kind(self, kind):\n    return replace(self, _memory_kind=kind)","handlingStrategy":"type-guard","validationCode":"assert type(s).with_memory_kind is not jax.sharding.Sharding.with_memory_kind, 'with_memory_kind not implemented'","typeGuard":"def has_with_memory_kind_impl(s) -> bool:\n    return type(s).with_memory_kind is not jax.sharding.Sharding.with_memory_kind","tryCatchPattern":null,"preventionTips":["Use built-in shardings for device_put(..., memory_kind=...) flows"],"tags":["jax","sharding","not-implemented","memory-kind"],"backgroundTag":"abstract-method-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}