{"record":{"id":"db4d9c4fa915bd76","repo":"jax-ml/jax","slug":"pointwise-operations-on-bitwidth-bit-types-are-u","errorCode":null,"errorMessage":"Pointwise operations on {bitwidth}-bit types are unsupported (except bitwise operations). Upcast to a 16- or 32-bit type before performing the operation.","messagePattern":"Pointwise operations on (.+?)-bit types are unsupported \\(except bitwise operations\\)\\. Upcast to a 16- or 32-bit type before performing the operation\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":1453,"sourceCode":"      return self.to_layout(WGMMA_LAYOUT_UPCAST_2X).to_layout(new_layout)\n    if not isinstance(self.layout, WGSplatFragLayout):\n      raise NotImplementedError(\n          f\"Cannot convert from {self.layout} to {new_layout}\"\n      )\n    return type(self).splat(\n        self.registers.item(), self.shape, new_layout, is_signed=self.is_signed\n    )\n\n  def _pointwise(\n      self,\n      op,\n      *other,\n      output_is_signed: bool | None = None,\n      restrict_bitwidth: bool = True,\n  ) -> FragmentedArray:\n    if restrict_bitwidth:\n      if (bitwidth := utils.bitwidth(self.mlir_dtype)) <= 8 and bitwidth != 1:\n        raise NotImplementedError(\n            f\"Pointwise operations on {bitwidth}-bit types are unsupported\"\n            \" (except bitwise operations). Upcast to a 16- or 32-bit type\"\n            \" before performing the operation.\"\n        )\n    # If our layout is a splat, then we should either dispatch to a non-splat\n    # layout, or broadcast ourselves to the output shape first.\n    if isinstance(self.layout, WGSplatFragLayout):\n      output_shape = self.shape\n      for i, o in enumerate(other):\n        if not isinstance(o, FragmentedArray):\n          continue\n        elif not isinstance(o.layout, WGSplatFragLayout):\n          return o._pointwise(\n              lambda o, this, *args: op(this, *args[:i], o, *args[i:]),\n              self,\n              *other[:i],\n              *other[i + 1 :],\n              output_is_signed=output_is_signed,","sourceCodeStart":1435,"sourceCodeEnd":1471,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L1435-L1471","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Upcast to i16/i32 (or f16/f32) before the arithmetic, then optionally downcast after","Use bitwise operations (and/or/xor/shifts) which are allowed on narrow types","Pass restrict_bitwidth=False only if you know the backend handles the op for that width (rare; internal)"],"exampleFix":"# before\nscaled = w4 * scale  # w4 is i4\n# after\nscaled = w4.cast(i32) * scale","handlingStrategy":"type-guard","validationCode":"bw = utils.bitwidth(x.mlir_dtype)\nif bw <= 8 and bw != 1:\n    x = x.cast(i32)  # upcast before arithmetic","typeGuard":"def pointwise_safe(x) -> bool:\n    bw = utils.bitwidth(x.mlir_dtype)\n    return bw == 1 or bw >= 16","tryCatchPattern":"try:\n    y = x + other\nexcept NotImplementedError as e:\n    if 'Pointwise operations' in str(e):\n        y = x.cast(i32) + other.cast(i32)\n    else:\n        raise","preventionTips":["Always upcast 4/8-bit quantized values before dequantization math","Reserve narrow types for bitwise ops and WGMMA operands only"],"tags":["jax","mosaic-gpu","pointwise","quantization","int4","int8","dtype-bitwidth"],"backgroundTag":"narrow-integer-arithmetic-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}