jax-ml/jax · error · NotImplementedError

{vec_bitwidth}

Error message

{vec_bitwidth}

What it means

Inside _lift_fast_packed_instr, the packed PTX instruction constraint character is only defined for 32-bit ('r') and 16-bit ('h') element bitwidths; other vector bitwidths (e.g. 64, 8) raise NotImplementedError(vec_bitwidth).

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:2007

  def _lift_fast_packed_instr(
      packed_instr: str, single_instr: str,
  ) -> Callable[[ir.Value, ir.Value], ir.Value]:
    def fast_instr(*args):
      arg_ty = original_arg_ty = args[0].type
      assert all(a.type == arg_ty for a in args)
      if not isinstance(arg_ty, ir.VectorType):
        args = [vector.broadcast(ir.VectorType.get((1,), arg_ty), a) for a in args]
      arg_ty = ir.VectorType(args[0].type)
      [vec_len] = arg_ty.shape
      vec_bitwidth = vec_len * utils.bitwidth(arg_ty.element_type)
      if vec_len == 1 or vec_bitwidth == 32:
        assert vec_bitwidth.bit_count() == 1
        if vec_bitwidth == 32:
          cstr = "r"
        elif vec_bitwidth == 16:
          cstr = "h"
        else:
          raise NotImplementedError(vec_bitwidth)
        int_ty = ir.IntegerType.get_signless(vec_bitwidth)
        args_ptx = ", ".join(f"${i}" for i in range(len(args) + 1))
        args_int = [utils.bitcast(a, int_ty) for a in args]
        result_int = llvm.inline_asm(
            int_ty,
            args_int,
            f"{single_instr if vec_len == 1 else packed_instr} {args_ptx};",
            f"={cstr}" + f",{cstr}" * len(args)
        )
        assert isinstance(result_int, ir.Value)
        return utils.bitcast(result_int, original_arg_ty)
      else:
        assert vec_bitwidth > 32
        slice_len = 32 // utils.bitwidth(arg_ty.element_type)
        offset = 0
        slices = []
        while offset < vec_len:
          slice_end = min(offset + slice_len, vec_len)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the non-approx, non-packed op path for unsupported widths
2. Compute in f32 (or f16) so the packed path applies
3. Report the missing width upstream if needed

Example fix

// before
y = x.tanh(approx=True)  # f64 vector fragment
// after
y = x.tanh()

w = utils.bitwidth(fa.mlir_dtype)
use_packed = w in (16, 32)

try:
    y = fast_op(x)
except NotImplementedError:
    y = precise_op(x)

When it happens

Trigger: Using packed fast instructions (e.g. packed approx math on vector registers) where each element's bitwidth is not 16 or 32, such as f64 vectors or 8-bit elements.

Common situations: Running approx/packed fast-math paths on f64 or sub-16-bit data after vectorization changes element widths.

Related errors

• Unsupported register bitwidth: {reg_bitwidth}
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}
• Unsupported memory space: {orig_ref_ty.memory_space}
• Unsupported memory space: {ref_type.memory_space}

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