jax-ml/jax · error · ValueError

Can not bitcast memory region of size {shape_bitwidth} bits

Error message

Can not bitcast memory region of size {shape_bitwidth} bits to dtype with {target_bitwidth} bits.

What it means

After bitcasting, the total size in bits of the memory region (num bytes * 8) must be evenly divisible by the target dtype's bitwidth, so the lowering can compute shape = shape_bits / target_bits. If the byte count doesn't divide evenly (e.g. 6 bytes to f32), the reshape is impossible and ValueError is raised.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:1501

  ref_ty = ir.MemRefType(ref.type)
  if not mgpu_utils.is_smem_ref(ref_ty):
    raise ValueError(f"Only workgroup memory is supported but got {ref}.")
  if len(ref_ty.shape) != 1:
    raise NotImplementedError(
        "Data type bitcast is only supported for 1D arrays."
    )
  [stride], _ = ref_ty.get_strides_and_offset()
  if stride != 1:
    raise ValueError(
        "Data type bitcast is only supported for contiguous 1D arrays, but got "
        f"stride={stride}."
    )
  [shape_bytes] = ref_ty.shape
  shape_bitwidth = shape_bytes * 8
  target_bitwidth = mgpu_utils.bitwidth(dst_dtype)

  if shape_bitwidth % target_bitwidth:
    raise ValueError(
        f"Can not bitcast memory region of size {shape_bitwidth} bits to dtype "
        f"with {target_bitwidth} bits."
    )

  result_type = ir.MemRefType.get(
      shape=(shape_bitwidth // target_bitwidth,),
      element_type=dst_dtype,
      memory_space=ref_ty.memory_space,
  )

  # Do a memref_ptr/ptr_as_memref roundtrip instead of using `memref.view`,
  # which refuses to take in our source ref. This is because `memref.view` only
  # works on a super restricted set of `memref`s. E.g., it does not work if an
  # offset is specified, which can be the case for our SMEM refs.
  return mgpu_utils.ptr_as_memref(mgpu_utils.memref_ptr(ref), result_type)


@overload

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Round the buffer allocation up to a multiple of the target dtype size (pad and slice after view)
Choose a target dtype whose bitwidth divides the region size (e.g. view as i8/i16 instead of f32)
Assert/validate byte length divisibility in kernel setup

Example fix

# before
buf = alloc_smem((6,), jnp.uint8))
v = buf.view(jnp.float32)  # 48 bits % 32 != 0
# after
buf = alloc_smem((8,), jnp.uint8))
v = buf.view(jnp.float32)  # 64 bits / 32 = 2 elements

Defensive patterns

Strategy: validation

Validate before calling

n_bytes = int(buf.shape[0]) if buf.dtype == jnp.uint8 else None
target_bits = jnp.dtype(target_dtype).itemsize * 8
assert n_bytes is not None and (n_bytes * 8) % target_bits == 0, f'{n_bytes}B not divisible by {target_dtype}'

Type guard

def fits_bitcast(n_bytes: int, target: jnp.dtype) -> bool:
    return (n_bytes * 8) % (jnp.dtype(target).itemsize * 8) == 0

Prevention

Allocate buffers in multiples of the widest target dtype size
Validate sizes once at kernel-launch config time
Pad dynamically-sized buffers to dtype multiples

When it happens

Trigger: Viewing an i8 SMEM buffer whose length in bytes is not a multiple of the target dtype size — e.g. buf of length 6 viewed as f32 (4 bytes each), or odd-sized buffers viewed as i16/f32/f64.

Common situations: Dynamic block sizes that produce non-multiple buffer lengths; packing sub-byte types then reinterpreting at a width that doesn't divide the total.

Related errors

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