{"record":{"id":"e9f80d19776ddb97","repo":"jax-ml/jax","slug":"per-operand-dot-precision-unsupported","errorCode":null,"errorMessage":"Per-operand dot precision unsupported","messagePattern":"Per-operand dot precision unsupported","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2932,"sourceCode":"\n    acc = arith.constant(\n        red_type, ir.DenseElementsAttr.get_splat(red_type, val)\n    )\n    red = vector.multi_reduction(\n        ir.Attribute.parse(\"#vector.kind<add>\"),\n        arith.mulf(x, y),\n        acc,\n        [1]\n    )\n    return vector.shape_cast(out_type, red)\n\n  tpu_dot_dims = jax_dot_dims_to_tpu_dot_dot_dims(\n      dimension_numbers, lhs_aval.shape, rhs_aval.shape\n  )\n\n  if precision is not None:\n    if precision[0] != precision[1]:\n      raise NotImplementedError(\"Per-operand dot precision unsupported\")\n    precision = precision[0]\n  if precision is None or precision == lax.Precision.DEFAULT:\n    precision_attr = None  # That's the default in Mosaic.\n  elif precision == lax.Precision.HIGH:\n    precision_attr = ir.Attribute.parse(\n        \"#tpu.contract_precision<bf16x3>\"\n    )\n  elif precision == lax.Precision.HIGHEST:\n    precision_attr = ir.Attribute.parse(\n        \"#tpu.contract_precision<fp32>\"\n    )\n  else:\n    raise NotImplementedError(f\"Unsupported dot precision: {precision}\")\n  out_tile = arith.constant(\n      out_type, ir.DenseElementsAttr.get_splat(out_type, val)\n  )\n  # Contracting second minor is to transpose the lhs. Only try fusing if it's\n  # an implicit transpose.","sourceCodeStart":2914,"sourceCodeEnd":2950,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2914-L2950","documentation":"lax dot primitives accept per-operand precision tuples like (Precision.DEFAULT, Precision.HIGH), but Mosaic TPU lowering maps dot to a single TPU contract precision attribute, so both operands must request the same precision.","triggerScenarios":"lax.dot_general / jnp.matmul with precision=(lax.Precision.DEFAULT, lax.Precision.HIGH) (mismatched tuple) inside a Pallas Mosaic TPU kernel.","commonSituations":"Code tuned for TPU-via-XLA where mixed precision hints were accepted, then moved into a Pallas kernel.","solutions":["Pass a uniform precision: precision=lax.Precision.HIGH or a matching tuple (p, p)","Pass precision=None for the default"],"exampleFix":"// before\nout = lax.dot_general(a, b, ..., precision=(lax.Precision.DEFAULT, lax.Precision.HIGHEST))\n// after\nout = lax.dot_general(a, b, ..., precision=lax.Precision.HIGHEST)","handlingStrategy":"validation","validationCode":"if isinstance(precision, tuple):\n    assert precision[0] == precision[1], 'per-operand precision must match on Mosaic'\nprecision = precision[0] if isinstance(precision, tuple) else precision","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pass scalar precision values","Normalize precision tuples in shared conv/dot helpers"],"tags":["jax","pallas","tpu","dot-general","precision"],"backgroundTag":"per-operand-precision-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}