jax-ml/jax · error · ValueError

{memref_ty} {static_offset=} is not divisible by {packing=}`

Error message

{memref_ty} {static_offset=} is not divisible by {packing=}`

What it means

When loading a memref whose element bitwidth is below 8 bits (e.g. i1, i4), the code packs elements into bytes (packing = 8 // elem_bitwidth). A static memref offset not divisible by that packing factor cannot be byte-addressed, so it raises. Note the message has a stray trailing backtick — a cosmetic bug in the format string.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:1860

  address_space = get_memref_llvm_address_space(memref_ty)
  ptr_ty = llvm.PointerType.get(address_space)
  desc_ty_fields = [ptr_ty, ptr_ty, i64]
  if rank > 0:
    desc_ty_fields += [llvm.ArrayType.get(i64, rank)] * 2
  desc_ty = llvm.StructType.get_literal(desc_ty_fields)
  desc = builtin.unrealized_conversion_cast([desc_ty], [memref_arg])
  assert isinstance(desc, ir.Value)
  aligned_ptr = llvm.extractvalue(ptr_ty, desc, [1])
  offset_elems = llvm.extractvalue(i64, desc, [2])

  elem_bitwidth = bitwidth(memref_ty.element_type)
  if elem_bitwidth < 8:
    *_, static_offset = memref_ty.get_strides_and_offset()
    if static_offset != ir.ShapedType.get_dynamic_stride_or_offset():
      assert elem_bitwidth.bit_count() == 1
      packing = 8 // elem_bitwidth
      if static_offset % packing != 0:
        raise ValueError(
            f"{memref_ty} {static_offset=} is not divisible by {packing=}`"
        )
      offset_bytes = c(static_offset // packing, i64)
    else:
      offset_bits = llvm.mul(
          offset_elems,
          c(elem_bitwidth, i64),
          overflow_flags=llvm.IntegerOverflowFlags.none,
      )
      offset_bytes = llvm.udiv(offset_bits, c(8, i64))
  else:
    assert elem_bitwidth % 8 == 0
    offset_bytes = llvm.mul(
        offset_elems,
        c(elem_bitwidth // 8, i64),
        overflow_flags=llvm.IntegerOverflowFlags.none,
    )
  return llvm.inttoptr(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Adjust the memref layout/slice so the static offset is a multiple of 8 // elem_bitwidth elements (e.g. even offsets for i4, multiple of 8 for i1)
2. Use a byte-aligned base pointer and express the element offset via strides instead of the static memref offset
3. If the offset is genuinely dynamic, ensure it is marked as a dynamic stride/offset rather than a misaligned static one

Example fix

# before
memref<64xi4, strided<[1], offset: 3>>  # 3 % 2 != 0 -> raises
# after
memref<64xi4, strided<[1], offset: 4>>  # byte-aligned for i4

bw = elem_bitwidth(memref_ty.element_type)
if bw < 8:
    packing = 8 // bw
    _, static_offset = memref_ty.get_strides_and_offset()
    assert static_offset % packing == 0, f'offset {static_offset} misaligned for packing {packing}'

When it happens

Trigger: Calling the load/pointer-bitcast path with a sub-byte element type (i1 giving packing=8, i4 giving packing=2) where the memref's static offset is, say, 3 for i4 (3 % 2 != 0). Arises with offset注意力 with packed quantized or boolean tensors.

Common situations: Loading quantized (4-bit) or boolean data through a strided/offset view created by slicing or padding; offsetting a sub-byte-type memref by an odd number of elements; using layouts that produce non-multiple static offsets on Blackwell paths.

Related errors

• 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.
• Can only transfer integer bytes (shape={shape}, dtype={dtype

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