jax-ml/jax · error · NotImplementedError

Unsupported packing: {self.packing}

Error message

Unsupported packing: {self.packing}

What it means

In _debug_print, TMEM words are unpacked according to self.packing; only the unpacking path (packing < full) and the full-packing bitcast are implemented. Any other packing value (e.g. packing inconsistent with the dtype) raises NotImplementedError.

Source

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

  def _debug_print(self) -> None:
    i32 = ir.IntegerType.get_signless(32)
    num_cols = self.layout.cols_in_shape(self.shape, utils.bitwidth(self.dtype))
    lane = arith.remui(utils.thread_idx(), arith.constant(i32, utils.WARPGROUP_SIZE))
    for c in range(num_cols):
      ptr = _tmem_addr_to_ptr(arith.addi(self.address, arith.constant(i32, c)))
      i32_vec = ir.VectorType.get((1,), i32)
      vec_val = nvvm.tcgen05_ld(i32_vec, nvvm.Tcgen05LdStShape.SHAPE_32X32B, ptr)
      val = vector.extract(vec_val, [], [0])
      dtype_bitwidth = utils.bitwidth(self.dtype)
      full_packing = 32 // dtype_bitwidth
      if self.packing == 1:
        if dtype_bitwidth < 32:
          val = arith.trunci(ir.IntegerType.get_signless(dtype_bitwidth), val)
        val = utils.bitcast(val, self.dtype)
      elif self.packing == full_packing:
        val = utils.bitcast(val, ir.VectorType.get((full_packing,), self.dtype))
      else:
        raise NotImplementedError(f"Unsupported packing: {self.packing}")
      # TODO(apaszke): Make this print logical, not physical location.
      utils.debug_print(f"[{{}}, {c}]: {{}}", lane, val, uniform=False)


def _transfer_32xcols(
    base_addr: ir.Value,
    cols: int,
    atom_shape: tuple[int, int],
    tmem_packing: int,
    reg_packing: int,
) -> Iterator[tuple[ir.Value, int, int, slice]]:
  """Generates a sequence of parameters for a given TMEM read or write.

  Arguments:
    base_addr: The base address of the TMEM region.
    cols: The number of logical columns to transfer.
    atom_shape: The logical shape of the tile written by the warp in a single
      TMEM transfer.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix the packing used at allocation so it equals full_packing or a supported divisor
2. Print via a supported packing (store to a default-layout TMEM and debug_print that)
3. Avoid _debug_print for exotic configurations; inspect via SMEM round-trip instead

full = 32 // utils.bitwidth(tmem.dtype)
assert tmem.packing == full or tmem.packing < full, 'unsupported packing for debug_print'

When it happens

Trigger: Calling tmem._debug_print() on a TensorMem whose packing is neither less than full_packing nor equal to full_packing — typically a mis-constructed allocation.

Common situations: Debugging kernels with unusual packing values; packing computed as 32//bitwidth producing values outside handled cases for exotic dtypes.

Related errors

• Unsupported packing: {tmem_packing} for element type {elt_bi
• Only {reg_packing} packing supported for bitwidth {elt_bitwi
• Only {reg_packing} supported for element type {dtype}, but g
• 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/e3a28b4bd844d441. Report an issue: GitHub.