jax-ml/jax · error · ValueError

Attempting to store into allocation with key `{key}` that is

Error message

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

What it means

store_buffer_content looked up a memory key whose stored object is not a Buffer. Stores need a Buffer to check dtype/shape and update race-detection clocks.

Source

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

      thread: The thread writing into the buffer.
      increment_clock: Whether to increment the given thread's vector clock.
      logging_info: Information about the source of the store.

    Returns:
      - True if the store was entirely in bounds, False otherwise (i.e. if the
        store was at least partially out of bounds).
      - The shape and dtype of the full content array of the buffer.
      - The incremented vector clock for the given thread.
        None if race detection is not enabled or if `increment_clock` is False.
    """
    clock = None
    with self.lock:
      if self.detect_races and increment_clock:
        clock = self.incr_clock(thread, take_lock=False)

      buff = self.mem[key]
      if not isinstance(buff, Buffer):
        raise ValueError(
            f"Attempting to store into allocation with key `{key}` that is not"
            " a `Buffer`."
        )
      shape_and_dtype = ShapeAndDtype(buff.logical_shape, buff.dtype)

      assert buff.dtype == value.dtype  # TODO(jburnim): Catch this statically.

      try:
        buff[rnge] = value
        is_in_bounds = True
      except IndexError:
        # `buf` was accessed with `rnge` at least partially out of bounds.
        is_in_bounds = False

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Verify the key refers to a Buffer allocation before storing
2. Eliminate collisions between fixed-ID semaphore IDs and buffer keys
3. Check that store ops reference the memory returned by allocation, not semaphore handles

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

try:
    mgr.store(key, value, ...)
except ValueError as e:
    if 'not a `Buffer`' in str(e):
        raise KeyError(f'{key} is not a buffer allocation') from e
    raise

When it happens

Trigger: Calling SharedMemManager.store (or ops routed through it) with a key that maps to a semaphore/non-Buffer allocation.

Common situations: Same family as the get variant: key collisions between fixed-ID semaphores and buffer keys, or passing a semaphore key to a store operation in a custom kernel interpretation.

Related errors

• Attempting to deallocate allocation with key `{key}` that is
• Attempting to get contents of allocation with key `{key}` th
• 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/1272fab9e4046974. Report an issue: GitHub.