jax-ml/jax · error · ValueError

Stored array has dtype {value.mlir_dtype}, but TMEM has dtyp

Error message

Stored array has dtype {value.mlir_dtype}, but TMEM has dtype {self.dtype}

What it means

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.

Source

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

    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.")
    packing = self.packing
    has_default_layout = self.layout == tmem_default_layout(packing=packing)
    bitwidth = utils.bitwidth(self.dtype)
    is_at_least_16b = bitwidth in {16, 32}
    if value.layout == LAYOUT and has_default_layout and is_at_least_16b:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the TMEM with the dtype of the value you will store (typically the accumulator dtype)
2. Convert the FragmentedArray's registers to the TMEM dtype (arith.truncf/extf/bitcast) before storing
3. Keep a single dtype constant shared by allocation and computation

Example fix

// before
tmem = TensorMem.alloc(shape, dtype=f32)
tmem.store(f16_result)
// after
tmem = TensorMem.alloc(shape, dtype=ir.F16Type.get())
tmem.store(f16_result)

if value.mlir_dtype != tmem.dtype:
    raise ValueError('dtype mismatch before store')

When it happens

Trigger: Allocating TMEM with dtype=f32 while the computed FragmentedArray holds f16/bf16 values, or integer vs float mismatches.

Common situations: Mixed-precision kernels where mma outputs f32 but TMEM was allocated for the input dtype; changing dtypes in one place but not the other.

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 shape {value.shape}, but TMEM has shape {se

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