jax-ml/jax · error · ValueError

a and b must have the same element type, but got: {a_type.el

Error message

a and b must have the same element type, but got: {a_type.element_type} and {b_type.element_type}

What it means

The Pallas Triton dot requires both operands to have the same element type at the MLIR level; mixed input dtypes (e.g. f16 * f32 or bf16 * f32 tiles) are rejected with ValueError before the accumulator is allocated.

Source

Thrown at jax/_src/pallas/triton/lowering.py:2430

    # to round to zero, corrupting the ValueTable and segfaulting the compiler.
    #   M >= 16  (2 × instrM=8)
    #   N >=  8  (1 × instrN=8)
    #   K >= 16  (4 × instrK=4)
    errors = []
    if m < 16:
      errors.append(f"M={m} < 16")
    if n < 8:
      errors.append(f"N={n} < 8")
    if k < 16:
      errors.append(f"K={k} < 16")
    if errors:
      raise ValueError(
          f"float64 dot requires M>=16, N>=8, K>=16 per warp tile "
          f"(Triton MMAv2 m8n8k4 layout); got {', '.join(errors)}"
      )

  if a_type.element_type != b_type.element_type:
    raise ValueError(
        "a and b must have the same element type, but got:"
        f" {a_type.element_type} and {b_type.element_type}"
    )

  assert acc_dtype is not None
  acc = _zeros(ir.RankedTensorType.get([m, n], _dtype_to_ir_type(acc_dtype)))

  if precision in (
      lax.DotAlgorithmPreset.BF16_BF16_F32_X3,
      lax.DotAlgorithmPreset.BF16_BF16_F32_X6,
      lax.DotAlgorithmPreset.BF16_BF16_F32_X9,
  ):
    a_bf16 = _as_bf16(a)
    b_bf16 = _as_bf16(b)
    a_err0 = _sub(a, _as_f32(a_bf16))
    b_err0 = _sub(b, _as_f32(b_bf16))
    a_err0_bf16 = _as_bf16(a_err0)
    b_err0_bf16 = _as_bf16(b_err0)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast both operands to a common dtype before the dot: pl.dot(a.astype(jnp.float32), b)
Or cast the lower-precision operand up: b = b.astype(a.dtype) at kernel level
Check that BlockSpec dtypes for both operands match the kernel's expected inputs

Example fix

# before
acc = pl.dot(a_f32, w_bf16)

# after
acc = pl.dot(a_f32, w_bf16.astype(jnp.float32))

Defensive patterns

Strategy: type-guard

Validate before calling

assert a.dtype == b.dtype, f'dtype mismatch: {a.dtype} vs {b.dtype}'

Type guard

def same_dtype(a, b) -> bool:
    return a.dtype == b.dtype

Prevention

Normalize all kernel operands to one dtype at launch (e.g. tree.map(lambda x: x.astype(jnp.float32), args))
Keep weight and activation dtypes aligned in checkpoints and loaders

When it happens

Trigger: Passing blocks of differing dtypes to pl.dot, e.g. one operand loaded as float16 and the other as float32, without an explicit cast in the kernel body.

Common situations: Loading quantized/mixed-precision weights (bf16) with activations in f32 and calling dot directly; kernels that worked on TPU mosaic where implicit promotion occurred.

Related errors

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