{"record":{"id":"3b03626c1e757ce9","repo":"jax-ml/jax","slug":"sharding-with-memory-kind-is-not-allowed-please-u","errorCode":null,"errorMessage":"sharding with memory_kind is not allowed. Please use `jax.device_put` to transfer to different memory spaces. Got {sharding=}","messagePattern":"sharding with memory_kind is not allowed\\. Please use `jax\\.device_put` to transfer to different memory spaces\\. Got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":2304,"sourceCode":"    * Checking for len(spec)-ndim match\n    * Checking if the mesh is an AbstractMesh.\n  \"\"\"\n  ndim = len(shape)\n  if sharding is None:\n    return _empty_sharding(ndim)\n\n  out_s = _maybe_modify_sharding(sharding, ndim)\n  if len(out_s.spec) != ndim:\n    raise ValueError(\n        f\"Length of sharding.spec ({len(out_s.spec)}) must be equal to aval's\"\n        f\" ndim ({ndim}). Got sharding.spec {out_s.spec}, aval.ndim {ndim} and\"\n        f\" sharding {out_s}\")\n  if not isinstance(out_s.mesh, mesh_lib.AbstractMesh):\n    raise ValueError(\"Mesh of an aval must be an AbstractMesh. \"\n                     f\"Got {out_s.mesh} of type {type(out_s.mesh)}\")\n  _check_divisibility(out_s, shape)\n  if out_s.memory_kind is not None:\n    raise ValueError(\n        \"sharding with memory_kind is not allowed. Please use `jax.device_put`\"\n        f\" to transfer to different memory spaces. Got {sharding=}\")\n  return out_s\n\n\n@cache(max_size=4096,\n       trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value)\ndef get_mat(mat, mesh):\n  if mesh.empty:\n    assert mat.empty, mat\n    return mat\n\n  axis_env = get_axis_env()\n  in_axis_env = lambda i: axis_env.axis_exists(i) and i not in mesh._name_to_type\n  for i in it.chain(mat.varying, mat.unreduced, mat.reduced):\n    if in_axis_env(i):\n      continue\n    if mesh._name_to_type[i] != AxisType.Manual:","sourceCodeStart":2286,"sourceCodeEnd":2322,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2286-L2322","documentation":"Shardings attached to abstract values (avals) must not pin a memory_kind (e.g. 'device' vs 'pinned' HBM/DRAM spaces); memory placement is a runtime concern. JAX tells you to use jax.device_put with an explicit destination to move data between memory spaces instead.","triggerScenarios":"Passing a sharding created with memory_kind=... (e.g. GSPMDSharding or NamedSharding with memory kind, or jax.sharding with MemoryKind) into an aval-producing path such as jit in_tree/sharding specs or custom primitive avals.","commonSituations":"Multi-memory-tier setups (GPU pinned host memory, TPU PBUF); porting code that hardcoded memory_kind in shardings; version changes disallowing memory_kind on aval shardings.","solutions":["Remove memory_kind from the sharding used for avals; keep it only for device_put","Use jax.device_put(x, jax.sharding.MemoryKind('pinned')) or device_put(x, sharding, memory_kind) to target a memory space","Construct two shardings: one plain for avals, one with memory_kind for transfers"],"exampleFix":"// before\nsh = GSPMDSharding(devices, partitions, memory_kind='pinned')\njit(fn, in_sharding=sh)(x)  # memory_kind rejected on aval\n\n// after\nsh = GSPMDSharding(devices, partitions)\njit(fn, in_sharding=sh)(jax.device_put(x, jax.sharding.MemoryKind('pinned')))","handlingStrategy":"validation","validationCode":"if getattr(sharding, 'memory_kind', None) is not None:\n    raise ValueError('strip memory_kind before aval use; use device_put')","typeGuard":"def is_plain_sharding(s): return getattr(s, 'memory_kind', None) is None","tryCatchPattern":null,"preventionTips":["Keep memory_kind only on device_put calls","Audit shardings for memory_kind before passing to jit in_sharding"],"tags":["jax","sharding","memory-kind","device-put"],"backgroundTag":"invalid-sharding-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}