{"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":"Mosaic conv lowering accepts only None/DEFAULT, and HIGHEST (fp32 contract precision). Other lax.Precision values such as HIGH, which are fine for dot, are not implemented for convolutions.","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","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}