jax-ml/jax · error · ValueError

Can only transfer integer bytes (shape={dst_ty.shape}, dtype

Error message

Can only transfer integer bytes (shape={dst_ty.shape}, dtype={dst_ty.element_type})

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:975

    assert isinstance(barrier_ref_aval, state_types.AbstractRef)
    base_index = _get_barrier_base_index(
        barrier_ref_aval,
        barrier_transforms_treedef.unflatten(flat_barrier_transforms),
    )
    if base_index is not None:
      barrier = barrier[base_index]
  collective = None
  if collective_axes is not None:
    collective = tuple(
        lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, axis)
        for axis in collective_axes
    )

  is_leader_tracked_copy = collective and leader_tracked is not None
  dst_ty = ir.MemRefType(dst.type)
  bits = math.prod(dst_ty.shape) * mgpu.bitwidth(dst_ty.element_type)
  if bits % 8:
    raise ValueError(
        f"Can only transfer integer bytes (shape={dst_ty.shape},"
        f" dtype={dst_ty.element_type})"
    )
  bytes = bits // 8

  if is_leader_tracked_copy:
    # Leader receives the completion messages from both CTAs.
    bytes *= 2
    if len(collective) != 1:
      raise ValueError(
          f"Expected exactly one collective axis, got {collective_axes=}"
      )
    if math.prod(ctx.launch_ctx.cluster_size) != 2:
      raise NotImplementedError(
          "Partitioned loads only supported for clusters of size 2. Got"
          f" cluster size {ctx.launch_ctx.cluster_size}."
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad or reshape the block so total bits % 8 == 0
2. Use int8 storage and unpack in SMEM
3. Verify the SMEM buffer shape matches whole-byte multiples of the dtype

Example fix

# before
copy_gmem_to_smem(gmem_int4[(7,)], smem)
# after
copy_gmem_to_smem(gmem_int4[(8,)], smem)  # 8 * 4 bits = 4 bytes

import math
bits = math.prod(dst_shape) * dtype_bitwidth(dtype)
assert bits % 8 == 0, 'copy size not byte-aligned'

When it happens

Trigger: 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.

Common situations: Quantized (int4/nibble) data experiments; malformed buffer shapes from dynamic shapes that collapse to odd sizes.

Related errors

• Can only transfer integer bytes (shape={shape}, dtype={dtype
• Only byte-aligned shapes are supported. Got shape: {ref.dtyp
• No valid out swizzle{what}: minor dimension has {minor_dim_b
• Can not bitcast memory region of size {shape_bitwidth} bits
• Only byte-aligned bitcasts are supported.

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