jax-ml/jax · error · ValueError

Expected same element type, got {element_type} and {dst_ref_

Error message

Expected same element type, got {element_type} and {dst_ref_ty.element_type}

What it means

async_copy in Mosaic GPU requires the source and destination memrefs to have identical element types. The check compares the MemRefType element types of both refs and raises this ValueError if they differ, since the underlying copy is a bitwise transfer with no implicit conversion.

Source

Thrown at jax/experimental/mosaic/gpu/launch_context.py:1263

      will have its expect_tx incremented by the total size of the transfer
      across all blocks involved in the collective. Barriers supplied by other
      blocks will be ignored (even if `arrive` is True).
    - If `leader_tracked` is ``CopyPartition.REPLICATED``, all blocks load the same data
      into their SMEM but only the first block in the collective tracks
      progress via barrier arrivals. This uses the `cta_group::2` mode.
    """
    index = ir.IndexType.get()
    i8 = ir.IntegerType.get_signless(8)
    i16 = ir.IntegerType.get_signless(16)
    i32 = ir.IntegerType.get_signless(32)
    i64 = ir.IntegerType.get_signless(64)

    src_ref_ty = ir.MemRefType(src_ref.type)
    dst_ref_ty = ir.MemRefType(dst_ref.type)
    element_type = src_ref_ty.element_type
    element_bitwidth = utils.bitwidth(element_type)
    if element_type != dst_ref_ty.element_type:
      raise ValueError(
          f"Expected same element type, got {element_type} and"
          f" {dst_ref_ty.element_type}"
      )

    if isinstance(collective, gpu.Dimension):
      collective = (collective,)
    elif collective is None:
      collective = ()
    if not isinstance(gmem_transform, tuple):
      gmem_transform = (gmem_transform,)
    if not isinstance(gmem_slice, tuple):
      gmem_slice = (gmem_slice,)

    if reduction_op is not None:
      if implementation != AsyncCopyImplementation.TMA:
        raise ValueError("Only the TMA implementation supports reductions")
      if not _is_tma_reduction_op_supported(reduction_op, element_type):
        raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make both refs use the same element type: allocate the SMEM buffer with the same dtype as the GMEM tensor.
2. If a dtype conversion is needed, perform async_copy at the source dtype and convert explicitly afterwards with an elementwise op.
3. Double-check memref creation utils so the address_space/memshape helper uses the tensor's dtype.

Example fix

// before
smem = mgpu.mem_ref((128, 64), jnp.float32)
ctx.async_copy(gmem_f16_ref, smem, ...)
// after
smem = mgpu.mem_ref((128, 64), jnp.float16)
ctx.async_copy(gmem_f16_ref, smem, ...)
# convert later: out = smem[:].astype(jnp.float32)

def same_element_type(src_ref, dst_ref) -> bool:
    return (ir.MemRefType(src_ref.type).element_type
            == ir.MemRefType(dst_ref.type).element_type)

When it happens

Trigger: Calling async_copy with src_ref of dtype f32 and dst_ref (SMEM buffer) of dtype bf16/f16, or any mismatched pair such as f16 source into an f32 SMEM allocation.

Common situations: Allocating SMEM buffers with a default dtype that differs from the global tensor dtype; changing kernel dtype (e.g. bf16 mixed precision) without updating smem allocation; copying f32 weights into half accumulators.

Related errors

• The stored value has dtype {src.dtype}, but the target refer
• Accumulator dtype {inner.dtype} does not match value dtype {
• All arrays must have the same dtype, got {arr.mlir_dtype} at
• The collective size ({collective_size}) must divide the slic
• Reduction op {reduction_op} not supported by the TMA impleme

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