jax-ml/jax · error · nb::value_error

Got invalid memory kind: %s. Valid memory kinds are: %s

Error message

Got invalid memory kind: %s. Valid memory kinds are: %s

What it means

NamedSharding's constructor validates an explicitly passed memory_kind against the list of valid memory kinds ('device', 'pinned_host_memory', 'unpinned_host_memory', and backend-specific kinds). An unknown string raises this value error listing the valid options.

Source

Thrown at jaxlib/sharding.cc:156

      mesh_(std::move(mesh)),
      spec_(std::move(spec)),
      memory_kind_(std::move(memory_kind)),
      logical_device_ids_(std::move(logical_device_ids)) {
  nb::object idl = nb::object(mesh_.attr("_internal_device_list"));
  if (idl.is_none()) {
    internal_device_list_ = std::nullopt;
  } else {
    internal_device_list_ = nb::cast<nb_class_ptr<PyDeviceList>>(idl);
  }
  if (internal_device_list_) {
    memory_kind_ =
        CheckAndCanonicalizeMemoryKind(memory_kind_, *internal_device_list_);
  } else {
    if (!memory_kind_.is_none() &&
        (std::find(valid_memory_kinds.begin(), valid_memory_kinds.end(),
                   nb::cast<std::string_view>(memory_kind_)) ==
         valid_memory_kinds.end())) {
      throw nb::value_error(
          absl::StrCat("Got invalid memory kind: ",
                       nb::cast<std::string_view>(memory_kind_),
                       ". Valid memory kinds are: ",
                       absl::StrJoin(valid_memory_kinds, ", "))
              .c_str());
    }
  }

  // TODO(phawkins): this leaks a reference to the check_pspec function.
  // A better way to fix this would be to move PartitionSpec and this check into
  // C++.
  static xla::SafeStatic<nb::object> check_pspec_init;
  nb::object& check_pspec = check_pspec_init.Get([]() {
    nb::module_ si = nb::module_::import_("jax._src.named_sharding");
    return std::make_unique<nb::object>(si.attr("check_pspec"));
  });
  check_pspec(mesh_, spec_);
}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use 'device' instead of 'hbm' for device/accelerator memory
2. Omit memory_kind to use the default device memory
3. Pick from the valid list in the message: typically 'device', 'pinned_host_memory', 'unpinned_host_memory'

Example fix

# before
NamedSharding(dev, memory_kind='hbm')

# after
NamedSharding(dev, memory_kind='device')

VALID = {'device', 'pinned_host_memory', 'unpinned_host_memory'}
if memory_kind is not None and memory_kind not in VALID:
    raise ValueError(f'use one of {VALID}')

def is_valid_memory_kind(kind: str | None) -> bool:
    return kind is None or kind in {'device', 'pinned_host_memory', 'unpinned_host_memory'}

try:
    sh = jax.sharding.NamedSharding(dev, memory_kind=kind)
except ValueError:
    sh = jax.sharding.NamedSharding(dev, memory_kind='device')

When it happens

Trigger: Constructing jax.sharding.NamedSharding(device, memory_kind='...') with a string not in the valid set — most commonly the literal 'hbm', which is not a valid kind (the HBM kind is named 'device').

Common situations: Assuming 'hbm' is a valid memory kind; copy-pasting device memory names from hardware docs into memory_kind; backend-specific kinds used on a different backend.

Related errors

• sharding with memory_kind is not allowed. Please use `jax.de
• Subclasses should implement this method
• Could not find memory addressable by device %s. Device %s ca
• 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/9b2e4520bcc627fc. Report an issue: GitHub.