When at least one mesh axis of `pred` is in auto mode, calli

Error message

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.

What it means

jax.lax.linalg-style error checking via jax.experimental.error_check.set_error_if performs a reduction of the predicate across devices. If any mesh axis of pred is in auto (SPMD partitioning) mode, that reduction requires implicit cross-device communication, which can be surprisingly slow or deadlock-prone; JAX warns before doing it.

Source

Thrown at jax/_src/error_check.py:186

  traceback = traceback.as_python_traceback()
  assert isinstance(traceback, TracebackType)
  traceback = traceback_util.filter_traceback(traceback)
  assert isinstance(traceback, TracebackType)

  with _error_list_lock:
    new_error_code = np.uint32(len(_error_list))
    _error_list.append((msg, traceback))

  out_sharding = core.typeof(_error_storage.ref).sharding
  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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Convert the relevant mesh axes to explicit mode (e.g. use explicit sharding for the array fed to set_error_if) so the reduction is explicit.
Re-shard pred to a single device or replicated sharding before calling set_error_if so no implicit collectives are needed.
If implicit communication is acceptable, suppress/ignore the RuntimeWarning, but profile the collective cost.

Example fix

# before
# auto-mode mesh → implicit communication warning
err = set_error_if(pred)  # pred sharded on auto axis
# after
# explicit mesh axis, or gather pred first
from jax.sharding import NamedSharding
pred_replicated = jax.device_put(pred, jax.sharding.Replicated())
err = set_error_if(pred_replicated)

Defensive patterns

Strategy: validation

Validate before calling

from jax.experimental.mesh_utils import AxisType  # or jax._src.mesh as mesh_lib
# ensure no Auto axes before set_error_if:
def no_auto_axes(mesh):
    types = getattr(mesh, 'axis_types', None)
    return not types or AxisType.Auto not in types
assert no_auto_axes(mesh)

Try / catch

with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter('always', RuntimeWarning)
    err = set_error_if(pred)
implicit_comm = any('implicit communication' in str(i.message) for i in w)

Prevention

Design error checks on explicit-sharding meshes; keep checks off auto-partitioned axes.
Replicate the predicate before set_error_if to make collectives explicit and profileable.
Document which mesh axes are Auto so downstream checks avoid them.

When it happens

Trigger: Calling set_error_if (or wrappers like _set_error_if_nan / _set_error_if_divide_by_zero, e.g. in custom derivatives or checks on sharded arrays) on a multi-device mesh where in_sharding.mesh.axis_types contains AxisType.Auto.

Common situations: Distributed training with auto-partitioned (GSPMD) meshes; NaN/divide-by-zero checks added to sharded computations; switching a manually-sharded pipeline to auto mode and hitting the warning on error checks.

Related errors

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