{"record":{"id":"ac3a3595785e2702","repo":"jax-ml/jax","slug":"unsupported-conv-precision-precision","errorCode":null,"errorMessage":"Unsupported conv precision: {precision}","messagePattern":"Unsupported conv precision: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2993,"sourceCode":"      f\"#tpu.conv_dimension_numbers<{lhs_spec[0]}, {lhs_spec[1]}, \"\n      f\"{format_dims(lhs_spec[2:])}, {rhs_spec[1]}, {rhs_spec[0]}, \"\n      f\"{format_dims(rhs_spec[2:])}, {out_spec[0]}, {out_spec[1]}, \"\n      f\"{format_dims(out_spec[2:])}>\"\n  )\n  return ir.Attribute.parse(tpu_conv_dims_str)\n\n\ndef _parse_precision_attr(precision):\n  if precision is not None:\n    if isinstance(precision, tuple) and precision[0] != precision[1]:\n      raise NotImplementedError(\"Per-operand conv precision unsupported\")\n    precision = precision[0] if isinstance(precision, tuple) else precision\n  if precision is None or precision == lax.Precision.DEFAULT:\n    return None\n  elif precision == lax.Precision.HIGHEST:\n    return ir.Attribute.parse(\"#tpu.contract_precision<fp32>\")\n  else:\n    raise NotImplementedError(f\"Unsupported conv precision: {precision}\")\n\n\n@register_lowering_rule(lax.conv_general_dilated_p)\ndef _conv_general_dilated_lowering_rule(\n    ctx: LoweringRuleContext,\n    lhs,\n    rhs,\n    *,\n    window_strides,\n    padding,\n    lhs_dilation,\n    rhs_dilation,\n    dimension_numbers,\n    feature_group_count,\n    batch_group_count,\n    precision=None,\n    preferred_element_type=None,\n    **_,","sourceCodeStart":2975,"sourceCodeEnd":3011,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2975-L3011","documentation":"Raised in _parse_precision_attr for Pallas Mosaic TPU convolutions: per-operand differing precisions and precision values outside the supported set have no Mosaic convolution attribute mapping, so they raise NotImplementedError. The input at fault is the precision argument of the conv primitive.","triggerScenarios":"lax.conv_general_dilated with precision=lax.Precision.HIGH (or any value other than DEFAULT/HIGHEST) in a Pallas Mosaic TPU kernel.","commonSituations":"Reusing dot precision settings for convs; assuming HIGH is universally supported because dot accepts it.","solutions":["Use lax.Precision.HIGHEST or lax.Precision.DEFAULT for convolutions on Mosaic","Pass precision=None"],"exampleFix":"// before\nout = lax.conv_general_dilated(x, w, ..., precision=lax.Precision.HIGH)\n// after\nout = lax.conv_general_dilated(x, w, ..., precision=lax.Precision.HIGHEST)","handlingStrategy":"validation","validationCode":"import jax.lax as lax\nassert precision in (None, lax.Precision.DEFAULT, lax.Precision.HIGHEST)  # conv only supports these","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Remember HIGH is dot-only on Mosaic; conv accepts DEFAULT/HIGHEST","Test each precision setting on TPU before shipping"],"tags":["jax","pallas","tpu","convolution","precision"],"backgroundTag":"unsupported-precision-mode","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","contentChangedAt":null,"schemaVersion":2},"datasetVersion":"2026-09-23T08:17:48.524Z"}