jax-ml/jax · error · NotImplementedError

Only {reg_packing} supported for element type {dtype}, but g

Error message

Only {reg_packing} supported for element type {dtype}, but got TMEM packing of {tmem_packing}

What it means

In the native load helper (_load_32xcols_native), when reg_packing != 2 the TMEM packing must exactly equal the register packing; otherwise there is no lowering and NotImplementedError is raised. Only 16-bit (reg_packing==2) loads can adapt tmem_packing in 1..2.

Source

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

  return vector_regs


def _load_32xcols_native(
    base_addr, cols, dtype, tmem_packing, vector_length, reduce: LoadReduceOp | None
) -> tuple[np.ndarray, ir.Value | None]:
  i32 = ir.IntegerType.get_signless(32)
  vec_ty = ir.VectorType.get((vector_length,), dtype)
  reg_packing = 32 // utils.bitwidth(dtype)
  assert vector_length % reg_packing == 0
  load_shape = "32x32b"
  load_atom_shape = (32, reg_packing)
  if reg_packing == 2:
    assert 1 <= tmem_packing <= 2
    pack = tmem_packing == 1
  else:
    if tmem_packing != reg_packing:
      raise NotImplementedError(
          f"Only {reg_packing} supported for element type {dtype}, but got"
          f" TMEM packing of {tmem_packing}"
      )
    pack = False

  it = _transfer_32xcols(base_addr, cols, load_atom_shape, tmem_packing, reg_packing)
  c0 = arith.constant(i32, 0)
  c1 = arith.constant(i32, 1)
  regs = [None] * (cols // reg_packing)
  red_reg = None
  for addr_row_col, instr_num, lane_step, num_slice in it:
    assert lane_step == 0, lane_step
    instr_regs = _tmem_load(addr_row_col, load_shape, instr_num, pack, reduce, dtype)
    if reduce:
      *instr_regs, instr_red_reg = instr_regs
      instr_red_reg = utils.bitcast(instr_red_reg, dtype)
      if red_reg is None:
        red_reg = instr_red_reg

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set tmem_packing == reg_packing at allocation time
2. For 16-bit types use reg_packing=2, which tolerates tmem_packing 1 or 2
3. Choose a different register layout that takes the non-native load path

Example fix

// before
tmem = TensorMem.alloc(shape, dtype=f32, packing=2)
arr = tmem.load(native_tiled_layout)
// after
tmem = TensorMem.alloc(shape, dtype=f32, packing=1)

if reg_packing != 2:
    assert tmem_packing == reg_packing, 'native load requires equal packings'

When it happens

Trigger: Calling load through the native path where tmem_packing != reg_packing for non-16-bit element types — e.g. packing=2 TMEM with scalar 32-bit register loads.

Common situations: Mismatched packing between TMEM allocation and the register layout chosen for load; kernel refactors that changed one packing but not the other.

Related errors

• Unsupported packing: {self.packing}
• Unsupported packing: {tmem_packing} for element type {elt_bi
• Only {reg_packing} packing supported for bitwidth {elt_bitwi
• Sparse meta layout loads unsupported.
• Loading multiple row tiles

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