jax-ml/jax · error · NotImplementedError

Unsupported conv precision

Error message

Unsupported conv precision: {precision}

What it means

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.

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
• Grouped convolutions are not supported on Pallas Mosaic TPU…
• Per-operand dot precision unsupported
• Requires libtpu >= 0.1.0
• Unsigned integer dtype

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

Appendix: 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)