jax-ml/jax · error · NotImplementedError

Requires libtpu >= 0.1.0

Error message

Requires libtpu >= 0.1.0

What it means

The main conv lowering rule gates on ctx.is_libtpu_at_least('0.1.0'): older libtpu builds lack the ConvOp used by Mosaic, so convolution in a Pallas TPU kernel raises until libtpu is upgraded.

Source

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

@register_lowering_rule(tpu_primitives.conv_p)
def _conv_lowering_rule(
    ctx: LoweringRuleContext,
    *args,
    dimension_numbers,
    window_strides=None,
    padding=None,
    lhs_dilation=None,
    rhs_dilation=None,
    window_reversal=None,
    feature_group_count=1,
    batch_group_count=1,
    precision=None,
    **_,
):
  if not ctx.is_libtpu_at_least("0.1.0"):
    # When removing this, also remove the pyrefly ignore annotation for ConvOp
    # below.
    raise NotImplementedError("Requires libtpu >= 0.1.0")

  if feature_group_count != 1 or batch_group_count != 1:
    raise NotImplementedError(
        "Grouped convolutions are not supported on Pallas Mosaic TPU backend"
        " yet."
    )
  for aval in ctx.avals_in[:2]:
    if jnp.issubdtype(aval.dtype, jnp.unsignedinteger):
      raise NotImplementedError(
          f"Unsigned integer dtype {aval.dtype} is not supported for conv on"
          " the Pallas Mosaic TPU backend."
      )
  lhs, rhs = args[0], args[1]
  acc = args[2] if len(args) > 2 else None
  (aval_out,) = ctx.avals_out
  out_type = ctx.aval_to_ir_type(aval_out)
  if acc is None:
    assert isinstance(out_type, ir.ShapedType)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upgrade libtpu to >= 0.1.0 (usually via upgrading jax/tpu extras: pip install -U jax[tpu]
2. Pin to a jax release that bundles a compatible libtpu
3. Avoid conv inside Pallas kernels and use regular lax conv outside until upgraded

# before launching conv kernels
import jax.lib.xla_client as xc
# or rely on jax version: keep jax[tpu] >= version bundling libtpu 0.1.0

try:
    kernel_with_conv(...)
except NotImplementedError as e:
    if 'libtpu' in str(e):
        run_conv_outside_pallas(...)  # fallback via plain lax

When it happens

Trigger: Using lax.conv_general_dilated in a Pallas Mosaic TPU kernel with a libtpu version below 0.1.0 installed in the JAX environment.

Common situations: Pinned/older jax/libtpu versions in a container or CI image; nightly-vs-release mismatches between jax and libtpu.

Related errors

• Pallas TPU requires a recent libtpu version (at least 0.0.46
• Per-operand conv precision unsupported
• Unsupported conv precision: {precision}
• Grouped convolutions are not supported on Pallas Mosaic TPU
• Unsigned integer dtype {aval.dtype} is not supported for con

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