{"record":{"id":"331b2b801c3dbf91","repo":"jax-ml/jax","slug":"requires-libtpu-0-1-0","errorCode":null,"errorMessage":"Requires libtpu >= 0.1.0","messagePattern":"Requires libtpu >= 0\\.1\\.0","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3049,"sourceCode":"@register_lowering_rule(tpu_primitives.conv_p)\ndef _conv_lowering_rule(\n    ctx: LoweringRuleContext,\n    *args,\n    dimension_numbers,\n    window_strides=None,\n    padding=None,\n    lhs_dilation=None,\n    rhs_dilation=None,\n    window_reversal=None,\n    feature_group_count=1,\n    batch_group_count=1,\n    precision=None,\n    **_,\n):\n  if not ctx.is_libtpu_at_least(\"0.1.0\"):\n    # When removing this, also remove the pyrefly ignore annotation for ConvOp\n    # below.\n    raise NotImplementedError(\"Requires libtpu >= 0.1.0\")\n\n  if feature_group_count != 1 or batch_group_count != 1:\n    raise NotImplementedError(\n        \"Grouped convolutions are not supported on Pallas Mosaic TPU backend\"\n        \" yet.\"\n    )\n  for aval in ctx.avals_in[:2]:\n    if jnp.issubdtype(aval.dtype, jnp.unsignedinteger):\n      raise NotImplementedError(\n          f\"Unsigned integer dtype {aval.dtype} is not supported for conv on\"\n          \" the Pallas Mosaic TPU backend.\"\n      )\n  lhs, rhs = args[0], args[1]\n  acc = args[2] if len(args) > 2 else None\n  (aval_out,) = ctx.avals_out\n  out_type = ctx.aval_to_ir_type(aval_out)\n  if acc is None:\n    assert isinstance(out_type, ir.ShapedType)","sourceCodeStart":3031,"sourceCodeEnd":3067,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3031-L3067","documentation":"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.","triggerScenarios":"Using lax.conv_general_dilated in a Pallas Mosaic TPU kernel with a libtpu version below 0.1.0 installed in the JAX environment.","commonSituations":"Pinned/older jax/libtpu versions in a container or CI image; nightly-vs-release mismatches between jax and libtpu.","solutions":["Upgrade libtpu to >= 0.1.0 (usually via upgrading jax/tpu extras: pip install -U jax[tpu]","Pin to a jax release that bundles a compatible libtpu","Avoid conv inside Pallas kernels and use regular lax conv outside until upgraded"],"exampleFix":null,"handlingStrategy":"fallback","validationCode":"# before launching conv kernels\nimport jax.lib.xla_client as xc\n# or rely on jax version: keep jax[tpu] >= version bundling libtpu 0.1.0","typeGuard":null,"tryCatchPattern":"try:\n    kernel_with_conv(...)\nexcept NotImplementedError as e:\n    if 'libtpu' in str(e):\n        run_conv_outside_pallas(...)  # fallback via plain lax","preventionTips":["Pin recent jax[tpu] versions in requirements","Check libtpu version in CI before running Pallas conv tests"],"tags":["jax","pallas","tpu","libtpu","version","convolution"],"backgroundTag":"library-version-too-old","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}