jax-ml/jax · error · NotImplementedError

Per-operand conv precision unsupported

Error message

Per-operand conv precision unsupported

What it means

The convolution precision parser in Mosaic only supports a single contract precision, so a precision tuple with different values per operand (e.g. (DEFAULT, HIGHEST)) for conv_general_dilated is rejected.

Source

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

  rhs_spec = dimension_numbers.rhs_spec
  out_spec = dimension_numbers.out_spec

  def format_dims(dims):
    return "[" + ", ".join(str(d) for d in dims) + "]"

  tpu_conv_dims_str = (
      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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a uniform precision (scalar or matching tuple)
2. Pass precision=None

Example fix

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

if isinstance(precision, tuple):
    assert precision[0] == precision[1]
precision = precision[0] if isinstance(precision, tuple) else precision

When it happens

Trigger: lax.conv_general_dilated with precision=(lax.Precision.DEFAULT, lax.Precision.HIGH) inside a Pallas Mosaic TPU kernel.

Common situations: Conv layers with per-operand precision hints ported from XLA/TPU pipelining code into Pallas.

Related errors

• Unsupported conv precision: {precision}
• 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/8e3e4ba5386a1f7c. Report an issue: GitHub.