{"record":{"id":"97f4db7f7d7793be","repo":"jax-ml/jax","slug":"expected-same-element-type-got-element-type-and","errorCode":null,"errorMessage":"Expected same element type, got {element_type} and {dst_ref_ty.element_type}","messagePattern":"Expected same element type, got (.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":1263,"sourceCode":"      will have its expect_tx incremented by the total size of the transfer\n      across all blocks involved in the collective. Barriers supplied by other\n      blocks will be ignored (even if `arrive` is True).\n    - If `leader_tracked` is ``CopyPartition.REPLICATED``, all blocks load the same data\n      into their SMEM but only the first block in the collective tracks\n      progress via barrier arrivals. This uses the `cta_group::2` mode.\n    \"\"\"\n    index = ir.IndexType.get()\n    i8 = ir.IntegerType.get_signless(8)\n    i16 = ir.IntegerType.get_signless(16)\n    i32 = ir.IntegerType.get_signless(32)\n    i64 = ir.IntegerType.get_signless(64)\n\n    src_ref_ty = ir.MemRefType(src_ref.type)\n    dst_ref_ty = ir.MemRefType(dst_ref.type)\n    element_type = src_ref_ty.element_type\n    element_bitwidth = utils.bitwidth(element_type)\n    if element_type != dst_ref_ty.element_type:\n      raise ValueError(\n          f\"Expected same element type, got {element_type} and\"\n          f\" {dst_ref_ty.element_type}\"\n      )\n\n    if isinstance(collective, gpu.Dimension):\n      collective = (collective,)\n    elif collective is None:\n      collective = ()\n    if not isinstance(gmem_transform, tuple):\n      gmem_transform = (gmem_transform,)\n    if not isinstance(gmem_slice, tuple):\n      gmem_slice = (gmem_slice,)\n\n    if reduction_op is not None:\n      if implementation != AsyncCopyImplementation.TMA:\n        raise ValueError(\"Only the TMA implementation supports reductions\")\n      if not _is_tma_reduction_op_supported(reduction_op, element_type):\n        raise ValueError(","sourceCodeStart":1245,"sourceCodeEnd":1281,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L1245-L1281","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Make both refs use the same element type: allocate the SMEM buffer with the same dtype as the GMEM tensor.","If a dtype conversion is needed, perform async_copy at the source dtype and convert explicitly afterwards with an elementwise op.","Double-check memref creation utils so the address_space/memshape helper uses the tensor's dtype."],"exampleFix":"// before\nsmem = mgpu.mem_ref((128, 64), jnp.float32)\nctx.async_copy(gmem_f16_ref, smem, ...)\n// after\nsmem = mgpu.mem_ref((128, 64), jnp.float16)\nctx.async_copy(gmem_f16_ref, smem, ...)\n# convert later: out = smem[:].astype(jnp.float32)","handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def same_element_type(src_ref, dst_ref) -> bool:\n    return (ir.MemRefType(src_ref.type).element_type\n            == ir.MemRefType(dst_ref.type).element_type)","tryCatchPattern":null,"preventionTips":["Allocate SMEM buffers via a helper that takes the tensor's dtype.","Add an assert on element types before async_copy.","Do explicit .astype conversions after the copy, never rely on implicit casting."],"tags":["jax","mosaic-gpu","async-copy","dtype-mismatch"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}