{"record":{"id":"ec0c25ad9d643a58","repo":"jax-ml/jax","slug":"unsupported-packing-tmem-packing-for-element-ty","errorCode":null,"errorMessage":"Unsupported packing: {tmem_packing} for element type {elt_bitwidth}","messagePattern":"Unsupported packing: (.+?) for element type (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1659,"sourceCode":"      c0 = arith.constant(i32, 0)\n      c1 = arith.constant(i32, 1)\n      for idx, vreg in enumerate(vector_regs):\n        regs[2 * idx] = llvm.extractelement(vreg, c0)\n        regs[2 * idx + 1] = llvm.extractelement(vreg, c1)\n    else:\n      regs = [utils.bitcast(r, i32) for r in vector_regs]\n    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]","sourceCodeStart":1641,"sourceCodeEnd":1677,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1641-L1677","documentation":"In the native store helper (_store_32xcols_native), when registers hold vectors, tmem_packing==1 is only implemented for 16-bit elements. For any other element bitwidth, packing=1 combined with vector registers has no lowering, so it raises NotImplementedError.","triggerScenarios":"Storing to TMEM allocated with packing=1 when the element type is not 16-bit and not 32-bit (e.g. 8-bit elements), i.e. tmem_packing==1 and elt_bitwidth!=32 with vector registers.","commonSituations":"Sub-16-bit or non-standard-width kernels where the native store path is selected with packing 1; mismatch between TMEM allocation packing and register packing.","solutions":["Allocate TMEM with the natural packing for the dtype (tmem_packing == 32 // elt_bitwidth)","Use the non-native store path (value.layout == LAYOUT with default TMEM layout)","For 16-bit elements, keep tmem_packing in 1..2 as the code supports"],"exampleFix":"// before\ntmem = TensorMem.alloc(..., packing=1)  # 8-bit elements\n# after: use natural packing\npacking = 32 // utils.bitwidth(dtype)","handlingStrategy":"validation","validationCode":"bw = utils.bitwidth(dtype)\nif tmem_packing == 1 and bw not in (16, 32):\n    tmem_packing = 32 // bw  # natural packing","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always compute packing as 32 // bitwidth at allocation time","Avoid hand-specifying packing=1 for sub-32-bit types"],"tags":["mosaic","gpu","tcgen05","packing","not-implemented","jax"],"backgroundTag":"unsupported-packing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}