{"record":{"id":"33cdcd9b4bf450df","repo":"jax-ml/jax","slug":"when-at-least-one-mesh-axis-of-pred-is-in-auto-m","errorCode":null,"errorMessage":"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.","messagePattern":"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\\.","errorType":"console","errorClass":null,"httpStatus":null,"severity":"warning","filePath":"jax/_src/error_check.py","lineNumber":186,"sourceCode":"  traceback = traceback.as_python_traceback()\n  assert isinstance(traceback, TracebackType)\n  traceback = traceback_util.filter_traceback(traceback)\n  assert isinstance(traceback, TracebackType)\n\n  with _error_list_lock:\n    new_error_code = np.uint32(len(_error_list))\n    _error_list.append((msg, traceback))\n\n  out_sharding = core.typeof(_error_storage.ref).sharding\n  in_sharding: NamedSharding = core.typeof(pred).sharding\n\n  # Reduce `pred`.\n  if all(dim is None for dim in out_sharding.spec):  # single-device case.\n    pred = pred.any()\n  else:  # multi-device case.\n    has_auto_axes = mesh_lib.AxisType.Auto in in_sharding.mesh.axis_types\n    if has_auto_axes:  # auto mode.\n      warnings.warn(\n          \"When at least one mesh axis of `pred` is in auto mode, calling\"\n          \" `set_error_if` will cause implicit communication between devices.\"\n          \" To avoid this, consider converting the mesh axis in auto mode to\"\n          \" explicit mode.\",\n          RuntimeWarning,\n      )\n      pred = pred.any()  # reduce to a single scalar\n    else:  # explicit mode.\n      if out_sharding.mesh != in_sharding.mesh:\n        raise ValueError(\n            \"The error code state and the predicate must be on the same mesh, \"\n            f\"but got {out_sharding.mesh} and {in_sharding.mesh} respectively. \"\n            \"Please use `with error_checking_context()` to redefine the error \"\n            \"code state based on the mesh.\"\n        )\n      pred = shard_map.shard_map(\n          partial(jnp.any, keepdims=True),\n          mesh=out_sharding.mesh,","sourceCodeStart":168,"sourceCodeEnd":204,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/error_check.py#L168-L204","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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."],"exampleFix":"# before\n# auto-mode mesh → implicit communication warning\nerr = set_error_if(pred)  # pred sharded on auto axis\n# after\n# explicit mesh axis, or gather pred first\nfrom jax.sharding import NamedSharding\npred_replicated = jax.device_put(pred, jax.sharding.Replicated())\nerr = set_error_if(pred_replicated)","handlingStrategy":"validation","validationCode":"from jax.experimental.mesh_utils import AxisType  # or jax._src.mesh as mesh_lib\n# ensure no Auto axes before set_error_if:\ndef no_auto_axes(mesh):\n    types = getattr(mesh, 'axis_types', None)\n    return not types or AxisType.Auto not in types\nassert no_auto_axes(mesh)","typeGuard":null,"tryCatchPattern":"with warnings.catch_warnings(record=True) as w:\n    warnings.simplefilter('always', RuntimeWarning)\n    err = set_error_if(pred)\nimplicit_comm = any('implicit communication' in str(i.message) for i in w)","preventionTips":["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."],"tags":["jax","distributed","spmd","sharding","error-check"],"backgroundTag":"implicit-cross-device-communication","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}