{"record":{"id":"220e6e4af9f8e665","repo":"jax-ml/jax","slug":"a-and-b-must-have-the-same-element-type-but-got","errorCode":null,"errorMessage":"a and b must have the same element type, but got: {a_type.element_type} and {b_type.element_type}","messagePattern":"a and b must have the same element type, but got: (.+?) and (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":2430,"sourceCode":"    # to round to zero, corrupting the ValueTable and segfaulting the compiler.\n    #   M >= 16  (2 × instrM=8)\n    #   N >=  8  (1 × instrN=8)\n    #   K >= 16  (4 × instrK=4)\n    errors = []\n    if m < 16:\n      errors.append(f\"M={m} < 16\")\n    if n < 8:\n      errors.append(f\"N={n} < 8\")\n    if k < 16:\n      errors.append(f\"K={k} < 16\")\n    if errors:\n      raise ValueError(\n          f\"float64 dot requires M>=16, N>=8, K>=16 per warp tile \"\n          f\"(Triton MMAv2 m8n8k4 layout); got {', '.join(errors)}\"\n      )\n\n  if a_type.element_type != b_type.element_type:\n    raise ValueError(\n        \"a and b must have the same element type, but got:\"\n        f\" {a_type.element_type} and {b_type.element_type}\"\n    )\n\n  assert acc_dtype is not None\n  acc = _zeros(ir.RankedTensorType.get([m, n], _dtype_to_ir_type(acc_dtype)))\n\n  if precision in (\n      lax.DotAlgorithmPreset.BF16_BF16_F32_X3,\n      lax.DotAlgorithmPreset.BF16_BF16_F32_X6,\n      lax.DotAlgorithmPreset.BF16_BF16_F32_X9,\n  ):\n    a_bf16 = _as_bf16(a)\n    b_bf16 = _as_bf16(b)\n    a_err0 = _sub(a, _as_f32(a_bf16))\n    b_err0 = _sub(b, _as_f32(b_bf16))\n    a_err0_bf16 = _as_bf16(a_err0)\n    b_err0_bf16 = _as_bf16(b_err0)","sourceCodeStart":2412,"sourceCodeEnd":2448,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L2412-L2448","documentation":"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.","triggerScenarios":"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.","commonSituations":"Loading quantized/mixed-precision weights (bf16) with activations in f32 and calling dot directly; kernels that worked on TPU mosaic where implicit promotion occurred.","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"],"exampleFix":"# before\nacc = pl.dot(a_f32, w_bf16)\n\n# after\nacc = pl.dot(a_f32, w_bf16.astype(jnp.float32))","handlingStrategy":"type-guard","validationCode":"assert a.dtype == b.dtype, f'dtype mismatch: {a.dtype} vs {b.dtype}'","typeGuard":"def same_dtype(a, b) -> bool:\n    return a.dtype == b.dtype","tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","pallas","triton","dtype","matmul"],"backgroundTag":"dtype-mismatch-in-matmul","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}