jax-ml/jax · error · ValueError

Stored array has shape {value.shape}, but TMEM has shape {se

Error message

Stored array has shape {value.shape}, but TMEM has shape {self.shape}

What it means

TensorMem.store requires the stored FragmentedArray to have exactly the same logical shape as the TMEM allocation. A shape mismatch means the allocation does not cover the data (or covers more), so the store is rejected.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:1461

      # that we didn't forget to handle reduce in any of the cases.
      assert reduced_reg is None
      return result
    reduced_layout = layout.reduce((len(layout.base_tile_shape) - 1,))
    assert reduced_layout.vector_length == 1
    reduced_regs_shape = reduced_layout.registers_shape(self.shape[:-1])
    assert math.prod(reduced_regs_shape) == 1
    reduced_result = fa.FragmentedArray(
        _registers=np.asarray(reduced_reg, dtype=object).reshape(reduced_regs_shape),
        _layout=reduced_layout,
        _is_signed=is_signed,
    )
    return result, reduced_result

  def store(self, value: fa.FragmentedArray):
    if not isinstance(value, fa.FragmentedArray):
      raise TypeError(f"TMEM stores expect a FragmentedArray, got: {value}")
    if value.shape != self.shape:
      raise ValueError(
          f"Stored array has shape {value.shape}, but TMEM has shape"
          f" {self.shape}"
      )
    if value.mlir_dtype != self.dtype:
      raise ValueError(
          f"Stored array has dtype {value.mlir_dtype}, but TMEM has dtype"
          f" {self.dtype}"
      )
    if not isinstance(value.layout, fa.TiledLayout):
      raise TypeError(f"Stored array has layout {value.layout}, but TMEM stores expect a TiledLayout")
    # TODO(olechwierowicz): `sparse_meta_layout()` does not really describe the
    # actual TMEM layout of the result of `async_copy_sparse_smem_to_tmem`.
    # As a result storing through SMEM -> Reg -> TMEM is not equivalent to
    # SMEM -> TMEM. We raise in this case to prevent inconsistent behaviour.
    # This restriction can be lifted if `TiledLayout` supports multiple
    # vector dims.
    if self.layout == sparse_meta_layout():
      raise NotImplementedError("Sparse meta layout stores unsupported.")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make value.shape match self.shape exactly (slice or pad the FragmentedArray)
2. Re-allocate TMEM with the shape of the value you produce
3. Assert shapes match before the store in kernel scaffolding

Example fix

// before
tmem_128x64.store(arr_64x64)
// after
tmem = TensorMem.alloc(shape=arr.shape, ...)
tmem.store(arr)

if value.shape != tmem.shape:
    raise ValueError(f'{value.shape} vs {tmem.shape}')

When it happens

Trigger: Allocating TMEM of shape (128, N) but storing a FragmentedArray of shape (64, N) or (128, N*2), e.g. after slicing or reshaping intermediate results.

Common situations: Splitting a computation over TMEM halves and forgetting to slice the array; changing tile sizes without resizing TMEM allocations.

Related errors

• Sparse meta layout loads unsupported.
• Loading multiple row tiles
• Loads from TMEM layout {self.layout} to register layout {lay
• TMEM stores expect a FragmentedArray, got: {value}
• Stored array has dtype {value.mlir_dtype}, but TMEM has dtyp

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