jax-ml/jax · error · ValueError

sharding with memory_kind is not allowed. Please use `jax.de

Error message

sharding with memory_kind is not allowed. Please use `jax.device_put` to transfer to different memory spaces. Got {sharding=}

What it means

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.

Source

Thrown at jax/_src/core.py:2304

    * Checking for len(spec)-ndim match
    * Checking if the mesh is an AbstractMesh.
  """
  ndim = len(shape)
  if sharding is None:
    return _empty_sharding(ndim)

  out_s = _maybe_modify_sharding(sharding, ndim)
  if len(out_s.spec) != ndim:
    raise ValueError(
        f"Length of sharding.spec ({len(out_s.spec)}) must be equal to aval's"
        f" ndim ({ndim}). Got sharding.spec {out_s.spec}, aval.ndim {ndim} and"
        f" sharding {out_s}")
  if not isinstance(out_s.mesh, mesh_lib.AbstractMesh):
    raise ValueError("Mesh of an aval must be an AbstractMesh. "
                     f"Got {out_s.mesh} of type {type(out_s.mesh)}")
  _check_divisibility(out_s, shape)
  if out_s.memory_kind is not None:
    raise ValueError(
        "sharding with memory_kind is not allowed. Please use `jax.device_put`"
        f" to transfer to different memory spaces. Got {sharding=}")
  return out_s


@cache(max_size=4096,
       trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value)
def get_mat(mat, mesh):
  if mesh.empty:
    assert mat.empty, mat
    return mat

  axis_env = get_axis_env()
  in_axis_env = lambda i: axis_env.axis_exists(i) and i not in mesh._name_to_type
  for i in it.chain(mat.varying, mat.unreduced, mat.reduced):
    if in_axis_env(i):
      continue
    if mesh._name_to_type[i] != AxisType.Manual:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove memory_kind from the sharding used for avals; keep it only for device_put
2. Use jax.device_put(x, jax.sharding.MemoryKind('pinned')) or device_put(x, sharding, memory_kind) to target a memory space
3. Construct two shardings: one plain for avals, one with memory_kind for transfers

Example fix

// before
sh = GSPMDSharding(devices, partitions, memory_kind='pinned')
jit(fn, in_sharding=sh)(x)  # memory_kind rejected on aval

// after
sh = GSPMDSharding(devices, partitions)
jit(fn, in_sharding=sh)(jax.device_put(x, jax.sharding.MemoryKind('pinned')))

if getattr(sharding, 'memory_kind', None) is not None:
    raise ValueError('strip memory_kind before aval use; use device_put')

def is_plain_sharding(s): return getattr(s, 'memory_kind', None) is None

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Subclasses should implement this method
• Could not find memory addressable by device %s. Device %s ca
• Got invalid memory kind: %s. Valid memory kinds are: %s
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m

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