{"record":{"id":"225d0e9c2ee22311","repo":"jax-ml/jax","slug":"can-only-transfer-integer-bytes-shape-dst-ty-sha","errorCode":null,"errorMessage":"Can only transfer integer bytes (shape={dst_ty.shape}, dtype={dst_ty.element_type})","messagePattern":"Can only transfer integer bytes \\(shape=(.+?), dtype=(.+?)\\)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":975,"sourceCode":"    assert isinstance(barrier_ref_aval, state_types.AbstractRef)\n    base_index = _get_barrier_base_index(\n        barrier_ref_aval,\n        barrier_transforms_treedef.unflatten(flat_barrier_transforms),\n    )\n    if base_index is not None:\n      barrier = barrier[base_index]\n  collective = None\n  if collective_axes is not None:\n    collective = tuple(\n        lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, axis)\n        for axis in collective_axes\n    )\n\n  is_leader_tracked_copy = collective and leader_tracked is not None\n  dst_ty = ir.MemRefType(dst.type)\n  bits = math.prod(dst_ty.shape) * mgpu.bitwidth(dst_ty.element_type)\n  if bits % 8:\n    raise ValueError(\n        f\"Can only transfer integer bytes (shape={dst_ty.shape},\"\n        f\" dtype={dst_ty.element_type})\"\n    )\n  bytes = bits // 8\n\n  if is_leader_tracked_copy:\n    # Leader receives the completion messages from both CTAs.\n    bytes *= 2\n    if len(collective) != 1:\n      raise ValueError(\n          f\"Expected exactly one collective axis, got {collective_axes=}\"\n      )\n    if math.prod(ctx.launch_ctx.cluster_size) != 2:\n      raise NotImplementedError(\n          \"Partitioned loads only supported for clusters of size 2. Got\"\n          f\" cluster size {ctx.launch_ctx.cluster_size}.\"\n      )\n","sourceCodeStart":957,"sourceCodeEnd":993,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L957-L993","documentation":"In _copy_gmem_to_smem_lowering, the transfer size in bits (prod(dst shape) * bitwidth(element_type)) must be divisible by 8. Sub-byte element types or sizes producing fractional bytes raise ValueError.","triggerScenarios":"copy_gmem_to_smem into an SMEM buffer whose total size in bits isn't byte-aligned, e.g. odd counts of 4-bit elements.","commonSituations":"Quantized (int4/nibble) data experiments; malformed buffer shapes from dynamic shapes that collapse to odd sizes.","solutions":["Pad or reshape the block so total bits % 8 == 0","Use int8 storage and unpack in SMEM","Verify the SMEM buffer shape matches whole-byte multiples of the dtype"],"exampleFix":"# before\ncopy_gmem_to_smem(gmem_int4[(7,)], smem)\n# after\ncopy_gmem_to_smem(gmem_int4[(8,)], smem)  # 8 * 4 bits = 4 bytes","handlingStrategy":"validation","validationCode":"import math\nbits = math.prod(dst_shape) * dtype_bitwidth(dtype)\nassert bits % 8 == 0, 'copy size not byte-aligned'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pad sub-byte-typed blocks to whole bytes","Validate dynamic shapes before launching kernels"],"tags":["jax","pallas","mosaic-gpu","sub-byte-dtype","alignment","gmem-copy"],"backgroundTag":"byte-alignment-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}