jax-ml/jax · error · ValueError

Passing lowering_platforms via jax.export or jit(f).trace(*a

Error message

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}

What it means

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.

Source

Thrown at jax/_src/interpreters/pxla.py:1049

  unique_in_shardings = util.stable_unique(in_shardings[len(const_args):])
  unique_out_shardings = util.stable_unique(out_shardings)
  backend, device_assignment, num_devices = _get_and_check_device_assignment(
      it.chain(
          ((i, stages.MismatchType.ARG_SHARDING, None) for i in unique_in_shardings),
          ((c, stages.MismatchType.CONST_SHARDING, None) for c in unique_const_shardings),
          ((o, stages.MismatchType.OUT_SHARDING, None) for o in unique_out_shardings),
          ((js, stages.MismatchType.SHARDING_INSIDE_COMPUTATION, source_info)
           for js, source_info in unique_intermediate_shardings)),
      context_mesh)
  unique_intermediate_shardings = [js for js, _ in unique_intermediate_shardings]
  unique_in_shardings = unique_in_shardings | unique_const_shardings  # pyrefly: ignore[unsupported-operation]
  del unique_const_shardings

  prim_requires_devices = dispatch.jaxpr_has_prim_requiring_devices(jaxpr)

  if device_assignment is None:
    if lowering_platforms is None:
      raise ValueError(
          "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"
          f" mesh: {context_mesh}")
    if prim_requires_devices:
      raise ValueError(
          "AbstractMesh cannot be used when jaxpr contains primitives that"
          " require devices to be present during lowering.")

  # For device_assignment == 1, this doesn't matter.
  if device_assignment is not None and len(device_assignment) > 1:
    rep_gs = GSPMDSharding.get_replicated(device_assignment)
    in_shardings = tuple(
        rep_gs if (isinstance(s, UnspecifiedValue) and
                   aval is not core.abstract_token and aval.ndim == 0)
        else s for s, aval in zip(in_shardings, global_in_avals))

  for a in global_out_avals:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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')
2. Alternatively give at least one argument a concrete sharding (real Mesh over real devices) so a platform can be inferred
3. For export, check the current jax.export API signature for the platforms parameter name

Example fix

# before
lowered = jit(f).trace(*args).lower()  # AbstractMesh-only shardings

# after
lowered = jit(f).trace(*args).lower(lowering_platforms='cpu')

import jax
from jax.sharding import NamedSharding
all_abstract = all(
    isinstance(s, NamedSharding) and isinstance(s.mesh, jax.sharding.AbstractMesh)
    for s in in_shardings
)
if all_abstract:
    assert lowering_platforms is not None, \
        'pass lowering_platforms when only AbstractMesh shardings are used'

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• AbstractMesh should be of the same size across all shardings
• AbstractMesh size: {abstract_mesh.size} does not match the d
• multi-platform lowering for buffer_callback
• Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of
• No VJP is available

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