jax-ml/jax · error · NotImplementedError

Unsupported conv precision: {precision}

Error message

Unsupported conv precision: {precision}

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:2993

      f"#tpu.conv_dimension_numbers<{lhs_spec[0]}, {lhs_spec[1]}, "
      f"{format_dims(lhs_spec[2:])}, {rhs_spec[1]}, {rhs_spec[0]}, "
      f"{format_dims(rhs_spec[2:])}, {out_spec[0]}, {out_spec[1]}, "
      f"{format_dims(out_spec[2:])}>"
  )
  return ir.Attribute.parse(tpu_conv_dims_str)


def _parse_precision_attr(precision):
  if precision is not None:
    if isinstance(precision, tuple) and precision[0] != precision[1]:
      raise NotImplementedError("Per-operand conv precision unsupported")
    precision = precision[0] if isinstance(precision, tuple) else precision
  if precision is None or precision == lax.Precision.DEFAULT:
    return None
  elif precision == lax.Precision.HIGHEST:
    return ir.Attribute.parse("#tpu.contract_precision<fp32>")
  else:
    raise NotImplementedError(f"Unsupported conv precision: {precision}")


@register_lowering_rule(lax.conv_general_dilated_p)
def _conv_general_dilated_lowering_rule(
    ctx: LoweringRuleContext,
    lhs,
    rhs,
    *,
    window_strides,
    padding,
    lhs_dilation,
    rhs_dilation,
    dimension_numbers,
    feature_group_count,
    batch_group_count,
    precision=None,
    preferred_element_type=None,
    **_,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use lax.Precision.HIGHEST or lax.Precision.DEFAULT for convolutions on Mosaic
2. Pass precision=None

Example fix

// before
out = lax.conv_general_dilated(x, w, ..., precision=lax.Precision.HIGH)
// after
out = lax.conv_general_dilated(x, w, ..., precision=lax.Precision.HIGHEST)

import jax.lax as lax
assert precision in (None, lax.Precision.DEFAULT, lax.Precision.HIGHEST)  # conv only supports these

When it happens

Trigger: lax.conv_general_dilated with precision=lax.Precision.HIGH (or any value other than DEFAULT/HIGHEST) in a Pallas Mosaic TPU kernel.

Common situations: Reusing dot precision settings for convs; assuming HIGH is universally supported because dot accepts it.

Related errors

• Per-operand conv precision unsupported
• Per-operand dot precision unsupported
• Unsupported dot precision: {precision}
• Grouped convolutions are not supported on Pallas Mosaic TPU
• Requires libtpu >= 0.1.0

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