jax-ml/jax · error · ValueError

Logical shape {self.logical_shape} cannot be bigger than con

Error message

Logical shape {self.logical_shape} cannot be bigger than content shape {self.content.shape}.

What it means

A shared-memory Buffer in Mosaic interpret mode was constructed with a logical_shape dimension larger than the underlying content array's shape. The logical shape may shrink (view a subregion) but can never exceed the allocated content.

Source

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

  def __init__(
      self,
      content: np.ndarray,
      ref_count: int = 1,
      logical_shape: tuple[int, ...] | None = None,
  ):
    super().__init__()

    self._content = content
    self._ref_count = ref_count

    if logical_shape is None:
      self._logical_shape = tuple(self._content.shape)
    else:
      self._logical_shape = logical_shape

    for dim, ldim in zip(self.content.shape, self.logical_shape, strict=True):
      if ldim > dim:
        raise ValueError(
            f"Logical shape {self.logical_shape} cannot be bigger than content"
            f" shape {self.content.shape}."
        )

  @property
  def content(self) -> np.ndarray:
    return self._content

  @property
  def ref_count(self) -> int:
    return self._ref_count

  @property
  def logical_shape(self) -> tuple[int, ...]:
    return self._logical_shape

  def decrease_ref_count(self):
    # We should never decrease the `ref_count` to below zero.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check that every dim of logical_shape <= the corresponding dim of the allocation; fix the allocation size or the logical shape
2. Inspect the kernel's BlockMapping and shared memory allocations for transposed or wrong-shape blocks
3. If writing tests that build Buffers directly, allocate content with the max needed shape

Example fix

// before
buf = Buffer(content=np.zeros((16, 32)), logical_shape=(32, 32))
// after
buf = Buffer(content=np.zeros((32, 32)), logical_shape=(16, 32))

assert all(l <= c for l, c in zip(logical_shape, content.shape, strict=True)), 'logical_shape exceeds content shape'

When it happens

Trigger: Creating interpret shared_memory.Buffer(content, logical_shape=...) where some logical dimension > the corresponding content.shape dimension; typically from BlockMapping with block shapes exceeding allocated SMEM.

Common situations: Mismatch between the BlockMapping's logical block shape and the allocated shared-memory buffer size in a Mosaic kernel; off-by-one or transposed block dims; changes to memory space allocation after a refactor.

Related errors

• Attempting to deallocate allocation with key `{key}` that is
• Attempting to get contents of allocation with key `{key}` th
• Attempting to store into allocation with key `{key}` that is
• Vector clock size ({self.vector_clock_size}) must be greater
• Range {rnge} is entirely out of bounds for shape {self.shape

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