jax-ml/jax · error · NotImplementedError

Pointwise operations on {bitwidth}-bit types are unsupported

Error message

Pointwise operations on {bitwidth}-bit types are unsupported (except bitwise operations). Upcast to a 16- or 32-bit type before performing the operation.

What it means

Mosaic GPU's _pointwise lowers arithmetic elementwise ops directly onto vector registers; sub-8-bit integer types (2/4/8-bit, bitwidth != 1) have no arithmetic lowering in this path, so the op refuses rather than emitting wrong code. Only bitwise ops and predicates (bitwidth 1) are exempt.

Source

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

      return self.to_layout(WGMMA_LAYOUT_UPCAST_2X).to_layout(new_layout)
    if not isinstance(self.layout, WGSplatFragLayout):
      raise NotImplementedError(
          f"Cannot convert from {self.layout} to {new_layout}"
      )
    return type(self).splat(
        self.registers.item(), self.shape, new_layout, is_signed=self.is_signed
    )

  def _pointwise(
      self,
      op,
      *other,
      output_is_signed: bool | None = None,
      restrict_bitwidth: bool = True,
  ) -> FragmentedArray:
    if restrict_bitwidth:
      if (bitwidth := utils.bitwidth(self.mlir_dtype)) <= 8 and bitwidth != 1:
        raise NotImplementedError(
            f"Pointwise operations on {bitwidth}-bit types are unsupported"
            " (except bitwise operations). Upcast to a 16- or 32-bit type"
            " before performing the operation."
        )
    # If our layout is a splat, then we should either dispatch to a non-splat
    # layout, or broadcast ourselves to the output shape first.
    if isinstance(self.layout, WGSplatFragLayout):
      output_shape = self.shape
      for i, o in enumerate(other):
        if not isinstance(o, FragmentedArray):
          continue
        elif not isinstance(o.layout, WGSplatFragLayout):
          return o._pointwise(
              lambda o, this, *args: op(this, *args[:i], o, *args[i:]),
              self,
              *other[:i],
              *other[i + 1 :],
              output_is_signed=output_is_signed,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Upcast to i16/i32 (or f16/f32) before the arithmetic, then optionally downcast after
  2. Use bitwise operations (and/or/xor/shifts) which are allowed on narrow types
  3. Pass restrict_bitwidth=False only if you know the backend handles the op for that width (rare; internal)

Example fix

# before
scaled = w4 * scale  # w4 is i4
# after
scaled = w4.cast(i32) * scale
Defensive patterns

Strategy: type-guard

Validate before calling

bw = utils.bitwidth(x.mlir_dtype)
if bw <= 8 and bw != 1:
    x = x.cast(i32)  # upcast before arithmetic

Type guard

def pointwise_safe(x) -> bool:
    bw = utils.bitwidth(x.mlir_dtype)
    return bw == 1 or bw >= 16

Try / catch

try:
    y = x + other
except NotImplementedError as e:
    if 'Pointwise operations' in str(e):
        y = x.cast(i32) + other.cast(i32)
    else:
        raise

Prevention

When it happens

Trigger: Calling arithmetic FragmentedArray ops (__add__, __sub__, __mul__, __neg__, clz, etc.) on an array with an i2/i4/i8 (or sub-8-bit custom) dtype without restrict_bitwidth=False; typically after quantized WGMMA or upcast layouts that carry 4-bit values.

Common situations: Doing dequantization math (scale/offset add) on 4-bit weights before upcasting; passing quantized fragments straight into pointwise ops; migrating int8 kernels to int4 without adding an explicit upcast step.

Related errors

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