jax-ml/jax · error · NotImplementedError

Unsupported packing: {tmem_packing} for element type {elt_bi

Error message

Unsupported packing: {tmem_packing} for element type {elt_bitwidth}

What it means

In the native store helper (_store_32xcols_native), when registers hold vectors, tmem_packing==1 is only implemented for 16-bit elements. For any other element bitwidth, packing=1 combined with vector registers has no lowering, so it raises NotImplementedError.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:1659

      c0 = arith.constant(i32, 0)
      c1 = arith.constant(i32, 1)
      for idx, vreg in enumerate(vector_regs):
        regs[2 * idx] = llvm.extractelement(vreg, c0)
        regs[2 * idx + 1] = llvm.extractelement(vreg, c1)
    else:
      regs = [utils.bitcast(r, i32) for r in vector_regs]
    assert tmem_packing == 1
    unpack = False
  elif reg_packing == 2:
    assert vector_length == 2
    # In this case, registers are already packed into 32-bit registers.
    regs = [utils.bitcast(r, i32) for r in vector_regs]
    if elt_bitwidth == 16:
      assert 1 <= tmem_packing <= 2
      unpack = tmem_packing == 1
    else:
      if tmem_packing == 1 and elt_bitwidth != 32:
        raise NotImplementedError(
            f"Unsupported packing: {tmem_packing} for element type {elt_bitwidth}"
        )
      assert tmem_packing == 32 // elt_bitwidth
      unpack = False
  else:
    if tmem_packing != reg_packing:
      raise NotImplementedError(
          f"Only {reg_packing} packing supported for bitwidth {elt_bitwidth},"
          f" but got TMEM packing of {tmem_packing}"
      )
    assert utils.bitwidth(vec_ty) == 32
    regs = [utils.bitcast(r, i32) for r in vector_regs]
    unpack = False
  cols = len(regs) * reg_packing
  it = _transfer_32xcols(base_addr, cols, store_atom_shape, tmem_packing, reg_packing)
  for addr_row_col, instr_num, lane_step, num_slice in it:
    assert lane_step == 0
    regs_slice = regs[num_slice]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Allocate TMEM with the natural packing for the dtype (tmem_packing == 32 // elt_bitwidth)
Use the non-native store path (value.layout == LAYOUT with default TMEM layout)
For 16-bit elements, keep tmem_packing in 1..2 as the code supports

Example fix

// before
tmem = TensorMem.alloc(..., packing=1)  # 8-bit elements
# after: use natural packing
packing = 32 // utils.bitwidth(dtype)

Defensive patterns

Strategy: validation

Validate before calling

bw = utils.bitwidth(dtype)
if tmem_packing == 1 and bw not in (16, 32):
    tmem_packing = 32 // bw  # natural packing

Prevention

Always compute packing as 32 // bitwidth at allocation time
Avoid hand-specifying packing=1 for sub-32-bit types

When it happens

Trigger: Storing to TMEM allocated with packing=1 when the element type is not 16-bit and not 32-bit (e.g. 8-bit elements), i.e. tmem_packing==1 and elt_bitwidth!=32 with vector registers.

Common situations: Sub-16-bit or non-standard-width kernels where the native store path is selected with packing 1; mismatch between TMEM allocation packing and register packing.

Related errors

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