{"record":{"id":"f18934045247be8a","repo":"jax-ml/jax","slug":"only-reg-packing-packing-supported-for-bitwidth","errorCode":null,"errorMessage":"Only {reg_packing} packing supported for bitwidth {elt_bitwidth}, but got TMEM packing of {tmem_packing}","messagePattern":"Only (.+?) packing supported for bitwidth (.+?), but got TMEM packing of (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1666,"sourceCode":"    assert tmem_packing == 1\n    unpack = False\n  elif reg_packing == 2:\n    assert vector_length == 2\n    # In this case, registers are already packed into 32-bit registers.\n    regs = [utils.bitcast(r, i32) for r in vector_regs]\n    if elt_bitwidth == 16:\n      assert 1 <= tmem_packing <= 2\n      unpack = tmem_packing == 1\n    else:\n      if tmem_packing == 1 and elt_bitwidth != 32:\n        raise NotImplementedError(\n            f\"Unsupported packing: {tmem_packing} for element type {elt_bitwidth}\"\n        )\n      assert tmem_packing == 32 // elt_bitwidth\n      unpack = False\n  else:\n    if tmem_packing != reg_packing:\n      raise NotImplementedError(\n          f\"Only {reg_packing} packing supported for bitwidth {elt_bitwidth},\"\n          f\" but got TMEM packing of {tmem_packing}\"\n      )\n    assert utils.bitwidth(vec_ty) == 32\n    regs = [utils.bitcast(r, i32) for r in vector_regs]\n    unpack = False\n  cols = len(regs) * reg_packing\n  it = _transfer_32xcols(base_addr, cols, store_atom_shape, tmem_packing, reg_packing)\n  for addr_row_col, instr_num, lane_step, num_slice in it:\n    assert lane_step == 0\n    regs_slice = regs[num_slice]\n    _tmem_store(addr_row_col, \"32x32b\", instr_num, regs_slice, unpack)\n\n\ndef _load_32xcols(base_addr, cols, dtype, tmem_packing) -> np.ndarray:\n  i32 = ir.IntegerType.get_signless(32)\n  vec_ty = ir.VectorType.get((2,), dtype)\n  reg_packing = 32 // utils.bitwidth(dtype)","sourceCodeStart":1648,"sourceCodeEnd":1684,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1648-L1684","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\ntmem = TensorMem.alloc(shape, dtype=f32, packing=2)\ntmem.store(scalar_f32_array)\n// after\ntmem = TensorMem.alloc(shape, dtype=f32, packing=1)","handlingStrategy":"validation","validationCode":"if tmem_packing != reg_packing:\n    raise ValueError('packings must match for scalar-register native store')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep TMEM packing and register packing derived from one config value","Assert packing equality before store in kernel setup"],"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"}