{"record":{"id":"8aad64be99a6b47e","repo":"jax-ml/jax","slug":"passing-lowering-platforms-via-jax-export-or-jit-f","errorCode":null,"errorMessage":"Passing lowering_platforms via jax.export or jit(f).trace(*args).lower(lowering_platforms=...) is required when only AbstractMesh exists in a jitted computation. Got context mesh: {context_mesh}","messagePattern":"Passing lowering_platforms via jax\\.export or jit\\(f\\)\\.trace\\(\\*args\\)\\.lower\\(lowering_platforms=\\.\\.\\.\\) is required when only AbstractMesh exists in a jitted computation\\. Got context mesh: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/pxla.py","lineNumber":1049,"sourceCode":"  unique_in_shardings = util.stable_unique(in_shardings[len(const_args):])\n  unique_out_shardings = util.stable_unique(out_shardings)\n  backend, device_assignment, num_devices = _get_and_check_device_assignment(\n      it.chain(\n          ((i, stages.MismatchType.ARG_SHARDING, None) for i in unique_in_shardings),\n          ((c, stages.MismatchType.CONST_SHARDING, None) for c in unique_const_shardings),\n          ((o, stages.MismatchType.OUT_SHARDING, None) for o in unique_out_shardings),\n          ((js, stages.MismatchType.SHARDING_INSIDE_COMPUTATION, source_info)\n           for js, source_info in unique_intermediate_shardings)),\n      context_mesh)\n  unique_intermediate_shardings = [js for js, _ in unique_intermediate_shardings]\n  unique_in_shardings = unique_in_shardings | unique_const_shardings  # pyrefly: ignore[unsupported-operation]\n  del unique_const_shardings\n\n  prim_requires_devices = dispatch.jaxpr_has_prim_requiring_devices(jaxpr)\n\n  if device_assignment is None:\n    if lowering_platforms is None:\n      raise ValueError(\n          \"Passing lowering_platforms via jax.export or\"\n          \" jit(f).trace(*args).lower(lowering_platforms=...) is required when\"\n          \" only AbstractMesh exists in a jitted computation. Got context\"\n          f\" mesh: {context_mesh}\")\n    if prim_requires_devices:\n      raise ValueError(\n          \"AbstractMesh cannot be used when jaxpr contains primitives that\"\n          \" require devices to be present during lowering.\")\n\n  # For device_assignment == 1, this doesn't matter.\n  if device_assignment is not None and len(device_assignment) > 1:\n    rep_gs = GSPMDSharding.get_replicated(device_assignment)\n    in_shardings = tuple(\n        rep_gs if (isinstance(s, UnspecifiedValue) and\n                   aval is not core.abstract_token and aval.ndim == 0)\n        else s for s, aval in zip(in_shardings, global_in_avals))\n\n  for a in global_out_avals:","sourceCodeStart":1031,"sourceCodeEnd":1067,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/pxla.py#L1031-L1067","documentation":"When a jitted computation's shardings only involve an AbstractMesh (no concrete devices anywhere), JAX has no device to derive the target platform from, so it requires lowering_platforms to be passed explicitly via jax.export or jit(f).trace(...).lower(lowering_platforms=...). Otherwise lower_sharding_computation raises.","triggerScenarios":"Creating a jitted function whose in/out shardings are all NamedShardings over an AbstractMesh (platform-agnostic sharding), then lowering without the lowering_platforms argument — e.g. with jax.export.export or manual .trace().lower() for portability.","commonSituations":"Using the export/AOT workflow with AbstractMesh to produce platform-portable artifacts; a second stricter case in the same code also rejects prims that need concrete devices when no device_assignment exists.","solutions":["Pass lowering_platforms: e.g. jit(f).trace(*args).lower(lowering_platforms='cpu') or the corresponding jax.export argument with your target (e.g. 'tpu', 'cuda')","Alternatively give at least one argument a concrete sharding (real Mesh over real devices) so a platform can be inferred","For export, check the current jax.export API signature for the platforms parameter name"],"exampleFix":"# before\nlowered = jit(f).trace(*args).lower()  # AbstractMesh-only shardings\n\n# after\nlowered = jit(f).trace(*args).lower(lowering_platforms='cpu')","handlingStrategy":"validation","validationCode":"import jax\nfrom jax.sharding import NamedSharding\nall_abstract = all(\n    isinstance(s, NamedSharding) and isinstance(s.mesh, jax.sharding.AbstractMesh)\n    for s in in_shardings\n)\nif all_abstract:\n    assert lowering_platforms is not None, \\\n        'pass lowering_platforms when only AbstractMesh shardings are used'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always specify lowering_platforms in export/AOT pipelines using AbstractMesh","Or attach one concrete sharding so the platform can be inferred"],"tags":["jax","abstract-mesh","lowering-platforms","export","aot","pjit"],"backgroundTag":"missing-required-configuration-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}