jax-ml/jax · error · TypeError

TMEM stores expect a FragmentedArray, got: {value}

Error message

TMEM stores expect a FragmentedArray, got: {value}

What it means

TensorMem.store type-checks its argument: only FragmentedArray values can be stored to tensor memory. Passing anything else (numpy array, list, tensor, etc.) raises TypeError immediately.

Source

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

    if reduce is None:
      # The None assignments in the branches let us use the linter to ensure
      # 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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap your registers: fa.FragmentedArray(_registers=regs, _layout=layout, _is_signed=...)
2. Build the value via mosaic ops (arith/mma results are already FragmentedArray)
3. Check isinstance(value, fa.FragmentedArray) before storing

Example fix

// before
tmem.store(my_numpy_array)
// after
fa_arr = fa.FragmentedArray(_registers=regs, _layout=layout, _is_signed=True)
tmem.store(fa_arr)

def is_fragmented_array(v) -> bool:
    return isinstance(v, fa.FragmentedArray)

try:
    tmem.store(value)
except TypeError as e:
    if 'FragmentedArray' in str(e):
        value = fa.FragmentedArray(_registers=value, _layout=layout, _is_signed=True)
        tmem.store(value)
    else: raise

When it happens

Trigger: tmem.store(np_array) or tmem.store(some_jax_array) — anything not an instance of mosaic FragmentedArray.

Common situations: New users treating TMEM store like an array assignment; passing registers or a raw vector instead of wrapping in FragmentedArray.

Related errors

• Stored array has layout {value.layout}, but TMEM stores expe
• Sparse meta layout loads unsupported.
• Loading multiple row tiles
• Loads from TMEM layout {self.layout} to register layout {lay
• Stored array has shape {value.shape}, but TMEM has shape {se

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