jax-ml/jax · error · ValueError

Types must match, got {high.type} and {low.type}

Error message

Types must match, got {high.type} and {low.type}

What it means

prmt (byte permute) combines two 32-bit registers, so both operands must have the same type. If high.type != low.type the types cannot be interpreted uniformly and it raises before the bitcasts to i32.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:2048

    raise ValueError("Only 32-bit scalar types supported")
  if isinstance(x.type, ir.IntegerType):
    pass
  elif isinstance(x.type, ir.F32Type):
    if get_arch().major != 10:
      raise ValueError("F32 redux only supported on Blackwell GPUs")
  else:
    raise NotImplementedError(x.type)
  assert mask.type == i32
  extra_kwargs: dict[str, Any] = {}
  if kind == ReductionKind.FMAX or kind == ReductionKind.FMIN:
    extra_kwargs = dict(nan=True)
  return nvvm.redux_sync(x, kind, mask, **extra_kwargs)


def prmt(high: ir.Value, low: ir.Value, permutation: ir.Value):
  i32 = ir.IntegerType.get_signless(32)
  if (result_type := high.type) != low.type:
    raise ValueError(f"Types must match, got {high.type} and {low.type}")
  if high.type != i32:
    high = bitcast(high, i32)
  if low.type != i32:
    low = bitcast(low, i32)
  if permutation.type != i32:
    permutation = bitcast(permutation, i32)
  result = llvm.inline_asm(
      i32, [high, low, permutation], "prmt.b32 $0, $1, $2, $3;", "=r,r,r,r"
  )
  assert isinstance(result, ir.Value)
  return bitcast(result, result_type)


def bitcast(x: ir.Value, new_type: ir.Type):
  if x.type == new_type:
    return x
  if (x_bw := bitwidth(x.type)) != (new_bw := bitwidth(new_type)):
    raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Bitcast one operand so both share the same type before calling prmt (prmt itself bitcasts to i32 afterwards, but requires equal input types)
2. Ensure upstream producers emit a consistent 32-bit type
3. Add a debug print/assert on high.type and low.type near the call site

Example fix

# before
prmt(high=f32_val, low=i32_val, perm=perm)
# after
prmt(high=bitcast(f32_val, ir.IntegerType.get_signless(32)), low=i32_val, perm=perm)

assert high.type == low.type, f'prmt operand mismatch: {high.type} vs {low.type}'

When it happens

Trigger: Calling prmt(high, low, perm) with mixed operand types, e.g. high of type i32 and low of type f32, or vector<2xi16> vs i32.

Common situations: Assembling a permute from values produced by different pipeline stages (one bitcast, one raw); refactoring shared code where operand dtypes drifted apart.

Related errors

• Only support memref.cast where the input and output types ar
• {op} has {name} of an unexpected type: {result}
• Expected an index-typed index
• Expected a dynamic index or an integer, got {offset}
• Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of

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