jax-ml/jax · error · ValueError

Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtyp

Error message

Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtype}.

What it means

During inline_mgpu lowering under lane (non-warp-group) semantics, each FragmentedArray value is type-checked against its declared ShapeDtypeStruct; the element type of the runtime MLIR value must equal the declared JAX dtype converted to an MLIR type, otherwise this dtype mismatch error is raised.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:3529

          if isinstance(r, state.AbstractRef)
      ),
  }


@discharge.register_discharge_rule(inline_mgpu_p)
def _inline_mgpu_discharge(*args, **kwargs):
  del args, kwargs
  raise NotImplementedError("inline_mgpu_p does not support discharge.")


def _type_check_mgpu_lane_semantics(v, ty):
  match (ty, v):
    case (RefType(), ir.Value()) if isinstance(v.type, ir.MemRefType):
      pass
    case (ShapeDtypeStruct(), mgpu.FragmentedArray()):
      mlir_dtype = mgpu_utils.dtype_to_ir_type(ty.dtype)
      if v.mlir_dtype != mlir_dtype:
        raise ValueError(
            f"Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtype}."
        )
      if ty.shape != v.shape:
        raise ValueError(
            f"Array shape mismatch: expected {ty.shape} got {v.shape}."
        )
      if v.layout != ty.layout.to_mgpu():
        raise ValueError(
            f"Array layout mismatch: expected {v.layout} got {ty.layout.to_mgpu()}."
        )
    case (SomeLayout(), mgpu.FragmentedArray()):
      if ty.to_mgpu() != v.layout:
        raise ValueError(f"Unexpected layout for {v} (expected: {ty})")
    case _:
      raise ValueError(f"Unexpected type {ty} for value {v}")


def _inline_mgpu_flat_transformed_args(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Align the declared dtype with what the MLIR op produces (cast inside the function or fix the annotation)
2. Insert an explicit cast (arith.extf/trunc or mgpu cast) inside the wrapped function

Example fix

# before
inline_mgpu(f, arg_types=[...], return_type=ShapeDtypeStruct(s, jnp.float32))
# after
inline_mgpu(f, arg_types=[...], return_type=ShapeDtypeStruct(s, jnp.bfloat16))  # match actual dtype

import jax.numpy as jnp
from jax._src.pallas.mosaic_gpu import mgpu_utils
assert v.mlir_dtype == mgpu_utils.dtype_to_ir_type(declared.dtype)

When it happens

Trigger: Declaring return_type/arg ShapeDtypeStruct with dtype jnp.float32 but the wrapped function produces an f64 or bf16 FragmentedArray.

Common situations: Implicit dtype promotion inside the inline function; declaring bf16 vs f16 incorrectly.

Related errors

• val.dtype={x.dtype} != ref.dtype={ref.dtype}
• The stored value has dtype {src.dtype}, but the target refer
• Accumulator dtype {inner.dtype} does not match value dtype {
• inline_mgpu_p only supports plgpu.ShapeDtypeStruct return ty
• inline_mgpu_p only supports only SomeLayout and RefType arg

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