{"record":{"id":"d02a0d7c8b389518","repo":"jax-ml/jax","slug":"the-error-code-state-and-the-predicate-must-be-on","errorCode":null,"errorMessage":"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.","messagePattern":"The error code state and the predicate must be on the same mesh, but got (.+?) and (.+?) respectively\\. Please use `with error_checking_context\\(\\)` to redefine the error code state based on the mesh\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/error_check.py","lineNumber":196,"sourceCode":"  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,\n          in_specs=in_sharding.spec,\n          out_specs=out_sharding.spec,\n      )(pred)  # perform per-device reduction\n\n  error_code = _error_storage.ref[...]\n  should_update = jnp.logical_and(error_code == jnp.uint32(_NO_ERROR), pred)\n  error_code = jnp.where(should_update, new_error_code, error_code)\n  # TODO(ayx): support vmap and shard_map.\n  _error_storage.ref[...] = error_code\n","sourceCodeStart":178,"sourceCodeEnd":214,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/error_check.py#L178-L214","documentation":"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().","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\npred = pjit_like_compute_on_mesh_b(...)  # different mesh\nset_error_if(pred)  # ValueError\n\n# after\nwith jax.error_check.error_checking_context(mesh_b):\n    set_error_if(pred)","handlingStrategy":"validation","validationCode":"from jax.experimental import error_checking as ec  # module name per version\nassert pred_sharding.mesh == error_state_mesh, 'mesh mismatch; re-enter error_checking_context'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","error-checking","sharding","mesh","distributed"],"backgroundTag":"mesh-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}