jax-ml/jax · error · NotImplementedError

Only packing supported for bitwidth , but got TMEM packing…

Error message

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

What it means

Raised by the TMEM unpacking path when the register packing implied by the layout is unsupported for the element bitwidth: the TMEM packing must equal the supported register packing (32/elt_bitwidth or 2). Any mismatch is rejected.

Solutions

1. Make tmem_packing equal reg_packing (usually 1 for 32-bit elements)
2. Repack registers into vectors matching the TMEM packing before storing
3. 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)

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

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

• Only supported for element type , but got TMEM packing of
• Unsupported packing
• Unsupported packing: for element type
• Loading multiple row tiles
• Sparse meta layout loads unsupported.

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

Appendix: 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)