jax-ml/jax · error · NotImplementedError

Only {reg_packing} packing supported for bitwidth {elt_bitwi

Error message

Only {reg_packing} packing supported for bitwidth {elt_bitwidth}, but got TMEM packing of {tmem_packing}

What it means

In _store_32xcols_native, for scalar (non-vector) registers the TMEM packing must equal the register packing; any mismatch has no lowering and raises NotImplementedError.

Source

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

    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]
    _tmem_store(addr_row_col, "32x32b", instr_num, regs_slice, unpack)


def _load_32xcols(base_addr, cols, dtype, tmem_packing) -> np.ndarray:
  i32 = ir.IntegerType.get_signless(32)
  vec_ty = ir.VectorType.get((2,), dtype)
  reg_packing = 32 // utils.bitwidth(dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make tmem_packing equal reg_packing (usually 1 for 32-bit elements)
Repack registers into vectors matching the TMEM packing before storing
Re-allocate TMEM with packing consistent with the register layout

Example fix

// before
tmem = TensorMem.alloc(shape, dtype=f32, packing=2)
tmem.store(scalar_f32_array)
// after
tmem = TensorMem.alloc(shape, dtype=f32, packing=1)

Defensive patterns

Strategy: validation

Validate before calling

if tmem_packing != reg_packing:
    raise ValueError('packings must match for scalar-register native store')

Prevention

Keep TMEM packing and register packing derived from one config value
Assert packing equality before store in kernel setup

When it happens

Trigger: Calling a store that routes to _store_32xcols_native where tmem_packing != reg_packing with 32-bit register elements — e.g. TMEM allocated with packing=2 while registers are scalar i32/f32.

Common situations: Allocating TMEM with a packing for sub-32-bit types but then storing full-width scalar registers; changing one of the two packings during a refactor.

Related errors

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