jax-ml/jax · error · ValueError

The error code state and the predicate must be on the same m

Error message

The error code state and the predicate must be on the same mesh, but got {out_sharding.mesh} and {in_sharding.mesh} respectively. Please use `with error_checking_context()` to redefine the error code state based on the mesh.

What it means

In explicit error-checking mode, set_error_if writes into a global error-code state array that is sharded on a specific mesh; the predicate array must live on the same mesh. If the predicate's sharding references a different mesh, JAX raises this ValueError telling you to rebuild the error state with error_checking_context().

Source

Thrown at jax/_src/error_check.py:196

  in_sharding: NamedSharding = core.typeof(pred).sharding

  # Reduce `pred`.
  if all(dim is None for dim in out_sharding.spec):  # single-device case.
    pred = pred.any()
  else:  # multi-device case.
    has_auto_axes = mesh_lib.AxisType.Auto in in_sharding.mesh.axis_types
    if has_auto_axes:  # auto mode.
      warnings.warn(
          "When at least one mesh axis of `pred` is in auto mode, calling"
          " `set_error_if` will cause implicit communication between devices."
          " To avoid this, consider converting the mesh axis in auto mode to"
          " explicit mode.",
          RuntimeWarning,
      )
      pred = pred.any()  # reduce to a single scalar
    else:  # explicit mode.
      if out_sharding.mesh != in_sharding.mesh:
        raise ValueError(
            "The error code state and the predicate must be on the same mesh, "
            f"but got {out_sharding.mesh} and {in_sharding.mesh} respectively. "
            "Please use `with error_checking_context()` to redefine the error "
            "code state based on the mesh."
        )
      pred = shard_map.shard_map(
          partial(jnp.any, keepdims=True),
          mesh=out_sharding.mesh,
          in_specs=in_sharding.spec,
          out_specs=out_sharding.spec,
      )(pred)  # perform per-device reduction

  error_code = _error_storage.ref[...]
  should_update = jnp.logical_and(error_code == jnp.uint32(_NO_ERROR), pred)
  error_code = jnp.where(should_update, new_error_code, error_code)
  # TODO(ayx): support vmap and shard_map.
  _error_storage.ref[...] = error_code

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Re-enter the context for the current mesh: with jax.error_check.error_checking_context(mesh): before the set_error_if calls
Ensure the predicate array is created/sharded on the same mesh as the error state (use the same sharding/NamedSharding)
Re-initialize error checking after any mesh change instead of reusing the prior state

Example fix

# before
pred = pjit_like_compute_on_mesh_b(...)  # different mesh
set_error_if(pred)  # ValueError

# after
with jax.error_check.error_checking_context(mesh_b):
    set_error_if(pred)

Defensive patterns

Strategy: validation

Validate before calling

from jax.experimental import error_checking as ec  # module name per version
assert pred_sharding.mesh == error_state_mesh, 'mesh mismatch; re-enter error_checking_context'

Prevention

Create error state and predicates under the same jax.make_mesh / sharding context
Re-run error_checking_context(mesh) whenever the mesh changes
Assert mesh equality in test harnesses for distributed error checking

When it happens

Trigger: Calling jax.error_check.set_error_if(pred) inside explicit-mode error checking where pred comes from an array sharded on mesh A while _error_storage.ref was initialized on mesh B — e.g. after re-sharding inputs or creating arrays under different jax.make_mesh contexts.

Common situations: Multi-host/multi-mesh training loops where the error state was set up before switching meshes; pipelines that reuse a global error state across model variants with different device meshes; refactors that moved set_error_if calls into functions operating on foreign-mesh arrays.

Related errors

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