jax-ml/jax · error · ValueError

trunc_div supports only integer types, got {self.mlir_dtype}

Error message

trunc_div supports only integer types, got {self.mlir_dtype}

What it means

FragmentedArray.trunc_div lowers to arith.divsi/divui, which only exist for integer types. If the array's mlir_dtype is a float (f16/bf16/f32), there is no signed/unsigned truncating division to emit, so it raises ValueError and asks you to use float division instead.

Source

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

      return self._e8m0_reciprocal()
    return self._pointwise(lambda s, o: arith.divf(o, s), other)

  def __floordiv__(self, other):
    if isinstance(self.mlir_dtype, ir.FloatType):
      return self._pointwise(
          lambda s, o: mlir_math.floor(arith.divf(s, o)), other
      )
    elif isinstance(self.mlir_dtype, ir.IntegerType):
      if self.is_signed:
        return self._pointwise(arith.floordivsi, other)
      else:
        return self._pointwise(arith.divui, other)
    else:
      return NotImplemented

  def trunc_div(self, other):
    if not isinstance(self.mlir_dtype, ir.IntegerType):
      raise ValueError(
          f"trunc_div supports only integer types, got {self.mlir_dtype}"
      )
    if self.is_signed:
      return self._pointwise(arith.divsi, other)
    else:
      return self._pointwise(arith.divui, other)

  def __rfloordiv__(self, other):
    if isinstance(self.mlir_dtype, ir.FloatType):
      return self._pointwise(
          lambda s, o: mlir_math.floor(arith.divf(o, s)), other
      )
    elif isinstance(self.mlir_dtype, ir.IntegerType):
      if self.is_signed:
        return self._pointwise(lambda s, o: arith.floordivsi(o, s), other)
      else:
        return self._pointwise(lambda s, o: arith.divui(o, s), other)
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use float division (__truediv__/arith.divf) for float dtypes
2. Branch on isinstance(arr.mlir_dtype, ir.IntegerType) before choosing trunc_div
3. If C-style truncation of floats is needed, compute divf then convert toward zero explicitly (e.g. floor for positive semantics)

Example fix

# before
q = a.trunc_div(b)  # a is f32
# after
q = a / b  # arith.divf for float types

from jax._src.lib import _mlir_dialects as d
assert isinstance(x.mlir_dtype, d.ir.IntegerType), "trunc_div needs an integer dtype"

def is_int_frag(x) -> bool:
    from jax._src.lib import _mlir_dialects as d
    return isinstance(x.mlir_dtype, d.ir.IntegerType)

try:
    q = x.trunc_div(y)
except ValueError as e:
    if 'trunc_div' in str(e):
        q = x / y
    else:
        raise

When it happens

Trigger: Calling trunc_div (or an API that dispatches to it, like _div for integer inputs) on a FragmentedArray whose dtype is a FloatType, e.g. arr.trunc_div(other) where arr is f32.

Common situations: Generic kernel code that picks trunc_div for 'round toward zero' semantics without checking dtype; reusing integer-kernel helper code with float operands; dtype parameterized by user config so integers work but float configs crash.

Related errors

• Denominator should be reduced along the same axes numerator
• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Cannot specify both out_shape and out_type. Use out_type.
• Cannot specify both scratch_shapes and scratch_types. Use sc

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