jax-ml/jax · error · NotImplementedError

Only byte-aligned bitcasts are supported.

Error message

Only byte-aligned bitcasts are supported.

What it means

When lowering an SMEM alias/bitcast, the total bit count of the transformed region (prod(shape) * bitwidth(dtype)) must be divisible by 8 because memory is byte-addressable. Sub-byte element counts (e.g. i1/i2/i4 blocks) that don't total whole bytes cannot be aliased and raise NotImplementedError.

Source

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

              f" dimension, got {ref.shape[0]} != {transformed_shape[0]}."
          )
        address = arith_dialect.addi(ref.address, _i32_constant(offset))
        ref = tcgen05.TMEMRef(
            address=address,
            shape=cast(tuple[int, int], transformed_shape),
            dtype=mlir_dtype,
            layout=layout,
        )
      else:
        assert isinstance(ref, ir.Value)  # make pyrefly happy
        input_ref_ty = ir.MemRefType(ref.type)
        if input_ref_ty.memory_space == mgpu_utils.smem():
          assert layout is None
          ref_bits = math.prod(transformed_shape) * mgpu_utils.bitwidth(
              mlir_dtype
          )
          if ref_bits % 8:
            raise NotImplementedError("Only byte-aligned bitcasts are supported.")
          assert offset % gpu_core.SMEM_ALIGNMENT == 0

          if lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
            if not isinstance(ref.owner, mgpu.dialect.SliceSMEMOp):
              # This restriction can be lifted by:
              # - Using memref ops to get the pointer and offset of the base ref.
              # - Subtracting gpu_dialect.dynamic_shared_memory() from those to
              #   get the base offset relative to the beginning of SMEM.
              # - Implementing layout and lowering rules for all ops above.
              raise NotImplementedError(
                  "The base ref for aliases must come from a slice_smem op."
              )

            base_offset = ref.owner.offset.value
            total_offset = base_offset + offset

            ref_ty = ir.MemRefType.get(
                transformed_shape, mlir_dtype, memory_space=mgpu_utils.smem()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad the block/element count so total bits are byte-aligned
2. Promote sub-byte dtypes to i8 before aliasing
3. Repack data so aliased regions start and end on byte boundaries

Example fix

# before
b = alloc_smem((3,), jnp.int4))  # 12 bits total
v = b.view(...)
# after
b = alloc_smem((4,), jnp.int4))  # 16 bits, byte-aligned
v = b.view(...)

total_bits = int(np.prod(shape)) * jnp.dtype(dt).itemsize * 8
assert total_bits % 8 == 0, f'{total_bits} bits is not byte-aligned'

def is_byte_aligned(shape, dt) -> bool:
    return (int(np.prod(shape)) * jnp.dtype(dt).itemsize * 8) % 8 == 0

When it happens

Trigger: Aliasing/biting an SMEM ref whose element count times element bitwidth is not a multiple of 8 — e.g. a block of 3 i4 values (12 bits), or boolean arrays of odd length.

Common situations: Using sub-byte packed dtypes (i4, i2, u1) in Pallas GPU kernels with view/alias operations; quantization kernels that pack low-bit weights.

Related errors

• Only workgroup memory is supported but got {ref}.
• Can not bitcast memory region of size {shape_bitwidth} bits
• packed, collective and layout arguments are only supported f
• Only byte-aligned shapes are supported. Got shape: {ref.dtyp
• All aliased Refs must have the same memory space (SMEM or TM

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