jax-ml/jax · error · NotImplementedError

{arg_ty}

Error message

{arg_ty}

What it means

Raised inside _lift_fast_instr's generated fast instruction: the register element type arg_ty is not a type the fast (PTX inline-asm) path supports for this operation, so no fast instruction can be emitted. This is an internal dispatch failure surfaced to the caller of approx pointwise ops.

Source

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

        for i in range(vec_len):
          vs = [
              vector.extract(
                  a,
                  dynamic_position=[],
                  static_position=ir.DenseI64ArrayAttr.get([i]),
              )
              for a in args
          ]
          vr = fast_instr(*vs)
          result = vector.insert(
              vr,
              result,
              dynamic_position=[],
              static_position=ir.DenseI64ArrayAttr.get([i]),
          )
        return result
      else:
        raise NotImplementedError(arg_ty)
    return fast_instr

  @staticmethod
  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"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Avoid approx=True for this dtype and use the full-precision op
2. Ensure the array is f32 before approx ops
3. If you believe the type should be supported, file an issue with jax.experimental.mosaic maintainers

Example fix

// before
y = x.sin(approx=True)  # unsupported register element type
// after
y = x.sin()

try_approx = isinstance(fa.mlir_dtype, ir.FloatType) and fa.mlir_dtype == ir.F32Type.get()

try:
    y = op(approx=True)
except NotImplementedError:
    y = op(approx=False)

When it happens

Trigger: Using _lift_fast_instr-wrapped ops (sin.approx/cos.approx/tanh.approx/rsqrt.approx, exp2 approx variants) on fragments whose register element type is not one of the supported float/int widths in the fast path.

Common situations: Calling approx math on unusual vector element types or after bitcasts that changed register element types; typically an internal Mosaic limitation rather than user error.

Related errors

• Unsupported reduction kind: {reduce}
• 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/299e74244cdfa23a. Report an issue: GitHub.