jax-ml/jax · error · ValueError

Attempting to deallocate allocation with key `{key}` that is

Error message

Attempting to deallocate allocation with key `{key}` that is not a `Buffer`.

What it means

The interpret-mode shared memory manager's deallocate_buffer looked up a memory key whose stored object is not a Buffer (e.g. a Semaphore). The ref-count-based deallocation path only applies to Buffers.

Source

Thrown at jax/_src/pallas/mosaic/interpret/shared_memory.py:480

        if self.enable_logging and logging_info is not None:
          self._log(
              logging_info.format(
                  f"{key=}, {ref_count=}.\nvalue_shape={value.shape},"
                  f" logical_shape={buff.logical_shape},"
                  f" content_shape={buff.shape}",
                  line_prefix="`allocate_buffer`",
              )
          )

  def deallocate_buffer(
      self, key: MemKey, logging_info: interpret_utils.LoggingInfo | None = None
  ):
    """Decreases the ref count for the buffer with `key` and deallocates the buffer if the ref count is zero."""
    with self.lock:
      buff = self.mem[key]
      if not isinstance(buff, Buffer):
        raise ValueError(
            f"Attempting to deallocate allocation with key `{key}` that is not"
            " a `Buffer`."
        )

      buff.decrease_ref_count()
      if buff.has_zero_ref_count():
        # TODO(paulbib): delete buffer from race detection state as well
        self.mem.pop(key)
        self.deallocated_bytes += buff.size
        del buff

        if self.enable_logging and logging_info is not None:
          self._log(
              logging_info.format(f"{key=}.", line_prefix="`deallocate_buffer`")
          )

      should_collect = self.deallocated_bytes > 100_000_000
      if should_collect:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Verify the key refers to an allocation created via allocate_buffer, not a semaphore
2. Avoid key collisions between fixed-ID semaphores and buffer keys
3. Audit custom alloc/dealloc call sites to keep keys symmetric

def is_buffer_key(mgr, key):
    with mgr.lock:
        return isinstance(mgr.mem.get(key), Buffer)

try:
    mgr.deallocate_buffer(key)
except ValueError as e:
    if 'not a `Buffer`' in str(e):
        pass  # key belongs to a semaphore; nothing to deallocate
    else:
        raise

When it happens

Trigger: Calling deallocate_buffer with a key that was registered for a semaphore or other non-Buffer allocation due to key collision between fixed-ID semaphores and buffer keys.

Common situations: Client-managed memory keys colliding with internal semaphore IDs; mixed allocation/deallocation bookkeeping bugs in custom interpret-mode tooling.

Related errors

• Attempting to get contents of allocation with key `{key}` th
• Attempting to store into allocation with key `{key}` that is
• Logical shape {self.logical_shape} cannot be bigger than con
• Attempting to swap into allocation with `key` {key} that is
• Vector clock size ({self.vector_clock_size}) must be greater

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