{"record":{"id":"191b3e9c56f66b1a","repo":"jax-ml/jax","slug":"only-reg-packing-supported-for-element-type-dty","errorCode":null,"errorMessage":"Only {reg_packing} supported for element type {dtype}, but got TMEM packing of {tmem_packing}","messagePattern":"Only (.+?) supported for element type (.+?), but got TMEM packing of (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1744,"sourceCode":"\n  return vector_regs\n\n\ndef _load_32xcols_native(\n    base_addr, cols, dtype, tmem_packing, vector_length, reduce: LoadReduceOp | None\n) -> tuple[np.ndarray, ir.Value | None]:\n  i32 = ir.IntegerType.get_signless(32)\n  vec_ty = ir.VectorType.get((vector_length,), dtype)\n  reg_packing = 32 // utils.bitwidth(dtype)\n  assert vector_length % reg_packing == 0\n  load_shape = \"32x32b\"\n  load_atom_shape = (32, reg_packing)\n  if reg_packing == 2:\n    assert 1 <= tmem_packing <= 2\n    pack = tmem_packing == 1\n  else:\n    if tmem_packing != reg_packing:\n      raise NotImplementedError(\n          f\"Only {reg_packing} supported for element type {dtype}, but got\"\n          f\" TMEM packing of {tmem_packing}\"\n      )\n    pack = False\n\n  it = _transfer_32xcols(base_addr, cols, load_atom_shape, tmem_packing, reg_packing)\n  c0 = arith.constant(i32, 0)\n  c1 = arith.constant(i32, 1)\n  regs = [None] * (cols // reg_packing)\n  red_reg = None\n  for addr_row_col, instr_num, lane_step, num_slice in it:\n    assert lane_step == 0, lane_step\n    instr_regs = _tmem_load(addr_row_col, load_shape, instr_num, pack, reduce, dtype)\n    if reduce:\n      *instr_regs, instr_red_reg = instr_regs\n      instr_red_reg = utils.bitcast(instr_red_reg, dtype)\n      if red_reg is None:\n        red_reg = instr_red_reg","sourceCodeStart":1726,"sourceCodeEnd":1762,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1726-L1762","documentation":"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.","triggerScenarios":"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.","commonSituations":"Mismatched packing between TMEM allocation and the register layout chosen for load; kernel refactors that changed one packing but not the other.","solutions":["Set tmem_packing == reg_packing at allocation time","For 16-bit types use reg_packing=2, which tolerates tmem_packing 1 or 2","Choose a different register layout that takes the non-native load path"],"exampleFix":"// before\ntmem = TensorMem.alloc(shape, dtype=f32, packing=2)\narr = tmem.load(native_tiled_layout)\n// after\ntmem = TensorMem.alloc(shape, dtype=f32, packing=1)","handlingStrategy":"validation","validationCode":"if reg_packing != 2:\n    assert tmem_packing == reg_packing, 'native load requires equal packings'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive both packings from the element bitwidth with one formula","For 16-bit types, use reg_packing=2 to get packing flexibility"],"tags":["mosaic","gpu","tcgen05","packing","not-implemented","jax"],"backgroundTag":"packing-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}