{"record":{"id":"e5585ecaf29b8893","repo":"jax-ml/jax","slug":"tmem-stores-expect-a-fragmentedarray-got-value","errorCode":null,"errorMessage":"TMEM stores expect a FragmentedArray, got: {value}","messagePattern":"TMEM stores expect a FragmentedArray, got: (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1459,"sourceCode":"    if reduce is None:\n      # The None assignments in the branches let us use the linter to ensure\n      # that we didn't forget to handle reduce in any of the cases.\n      assert reduced_reg is None\n      return result\n    reduced_layout = layout.reduce((len(layout.base_tile_shape) - 1,))\n    assert reduced_layout.vector_length == 1\n    reduced_regs_shape = reduced_layout.registers_shape(self.shape[:-1])\n    assert math.prod(reduced_regs_shape) == 1\n    reduced_result = fa.FragmentedArray(\n        _registers=np.asarray(reduced_reg, dtype=object).reshape(reduced_regs_shape),\n        _layout=reduced_layout,\n        _is_signed=is_signed,\n    )\n    return result, reduced_result\n\n  def store(self, value: fa.FragmentedArray):\n    if not isinstance(value, fa.FragmentedArray):\n      raise TypeError(f\"TMEM stores expect a FragmentedArray, got: {value}\")\n    if value.shape != self.shape:\n      raise ValueError(\n          f\"Stored array has shape {value.shape}, but TMEM has shape\"\n          f\" {self.shape}\"\n      )\n    if value.mlir_dtype != self.dtype:\n      raise ValueError(\n          f\"Stored array has dtype {value.mlir_dtype}, but TMEM has dtype\"\n          f\" {self.dtype}\"\n      )\n    if not isinstance(value.layout, fa.TiledLayout):\n      raise TypeError(f\"Stored array has layout {value.layout}, but TMEM stores expect a TiledLayout\")\n    # TODO(olechwierowicz): `sparse_meta_layout()` does not really describe the\n    # actual TMEM layout of the result of `async_copy_sparse_smem_to_tmem`.\n    # As a result storing through SMEM -> Reg -> TMEM is not equivalent to\n    # SMEM -> TMEM. We raise in this case to prevent inconsistent behaviour.\n    # This restriction can be lifted if `TiledLayout` supports multiple\n    # vector dims.","sourceCodeStart":1441,"sourceCodeEnd":1477,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1441-L1477","documentation":"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.","triggerScenarios":"tmem.store(np_array) or tmem.store(some_jax_array) — anything not an instance of mosaic FragmentedArray.","commonSituations":"New users treating TMEM store like an array assignment; passing registers or a raw vector instead of wrapping in FragmentedArray.","solutions":["Wrap your registers: fa.FragmentedArray(_registers=regs, _layout=layout, _is_signed=...)","Build the value via mosaic ops (arith/mma results are already FragmentedArray)","Check isinstance(value, fa.FragmentedArray) before storing"],"exampleFix":"// before\ntmem.store(my_numpy_array)\n// after\nfa_arr = fa.FragmentedArray(_registers=regs, _layout=layout, _is_signed=True)\ntmem.store(fa_arr)","handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def is_fragmented_array(v) -> bool:\n    return isinstance(v, fa.FragmentedArray)","tryCatchPattern":"try:\n    tmem.store(value)\nexcept TypeError as e:\n    if 'FragmentedArray' in str(e):\n        value = fa.FragmentedArray(_registers=value, _layout=layout, _is_signed=True)\n        tmem.store(value)\n    else: raise","preventionTips":["Always build values through mosaic ops or explicit FragmentedArray construction","Type-check at kernel boundaries in your wrappers"],"tags":["mosaic","gpu","tcgen05","type-check","tensor-memory","jax"],"backgroundTag":"wrong-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}