{"record":{"id":"3d3b35029771dc49","repo":"jax-ml/jax","slug":"stored-array-has-dtype-value-mlir-dtype-but-tme","errorCode":null,"errorMessage":"Stored array has dtype {value.mlir_dtype}, but TMEM has dtype {self.dtype}","messagePattern":"Stored array has dtype (.+?), but TMEM has dtype (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1466,"sourceCode":"    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.\n    if self.layout == sparse_meta_layout():\n      raise NotImplementedError(\"Sparse meta layout stores unsupported.\")\n    packing = self.packing\n    has_default_layout = self.layout == tmem_default_layout(packing=packing)\n    bitwidth = utils.bitwidth(self.dtype)\n    is_at_least_16b = bitwidth in {16, 32}\n    if value.layout == LAYOUT and has_default_layout and is_at_least_16b:","sourceCodeStart":1448,"sourceCodeEnd":1484,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1448-L1484","documentation":"TensorMem.store requires the stored array's MLIR dtype to equal the TMEM allocation's dtype exactly. Mismatched element types (e.g. f16 array into f32 TMEM, i32 into f32) are rejected rather than implicitly converted.","triggerScenarios":"Allocating TMEM with dtype=f32 while the computed FragmentedArray holds f16/bf16 values, or integer vs float mismatches.","commonSituations":"Mixed-precision kernels where mma outputs f32 but TMEM was allocated for the input dtype; changing dtypes in one place but not the other.","solutions":["Allocate the TMEM with the dtype of the value you will store (typically the accumulator dtype)","Convert the FragmentedArray's registers to the TMEM dtype (arith.truncf/extf/bitcast) before storing","Keep a single dtype constant shared by allocation and computation"],"exampleFix":"// before\ntmem = TensorMem.alloc(shape, dtype=f32)\ntmem.store(f16_result)\n// after\ntmem = TensorMem.alloc(shape, dtype=ir.F16Type.get())\ntmem.store(f16_result)","handlingStrategy":"validation","validationCode":"if value.mlir_dtype != tmem.dtype:\n    raise ValueError('dtype mismatch before store')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Allocate TMEM with the accumulator/output dtype, not the input dtype","Share one dtype constant between allocation and compute code"],"tags":["mosaic","gpu","tcgen05","dtype-mismatch","tensor-memory","jax"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}