{"record":{"id":"7ba4d05e65d4a6db","repo":"jax-ml/jax","slug":"mesh-must-be-provided-for-shard-map-with-checkify","errorCode":null,"errorMessage":"Mesh must be provided for shard_map with checkify.","messagePattern":"Mesh must be provided for shard_map with checkify\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/checkify.py","lineNumber":978,"sourceCode":"                                 *vals_in, _prim=_prim)\n  err_vals, err_tree = jtu.tree_flatten(error)\n  new_vals_in = [*err_vals, *vals_in]\n  in_avals = tuple(map(core.typeof, new_vals_in))\n  checked_jaxpr_, out_tree, _ = jaxpr_to_checkify_jaxpr(\n      _prim.jaxpr, enabled_errors, err_tree, *in_avals)\n  checked_jaxpr, consts = pe.separate_consts(checked_jaxpr_)\n  new_prim = ad_checkpoint.RematTraced(checked_jaxpr, _prim.policy)\n  err_and_out = new_prim(*consts, *new_vals_in)\n  return tree_unflatten(out_tree, err_and_out)\nerror_checks[call_hi_primitive_p] = call_hi_primitive_error_check\n\n\ndef shard_map_error_check(\n    error: Error, enabled_errors, *vals_in,\n    jaxpr: core.Jaxpr, in_specs, out_specs, **kwargs\n):\n  if (mesh := kwargs.get('mesh')) is None:\n    raise ValueError('Mesh must be provided for shard_map with checkify.')\n\n  err_vals, err_tree = jtu.tree_flatten(error)\n  num_error_vals = len(err_vals)\n  # Replicated sharding for in errors.\n  new_in_specs = (*([P()] * num_error_vals), *in_specs)\n  new_vals_in = [*err_vals, *vals_in]\n  in_avals = list(map(core.typeof, new_vals_in))\n  manual_axes = kwargs.get('newly_manual_axes')\n  check_vma = kwargs.get('check_vma')\n  for i, v in enumerate(in_avals):\n    if not (sharder := core.shard_aval_handlers.get(type(v))):\n      raise ValueError(f'Unsupported aval type: {type(v)}')\n    in_avals[i] = sharder(mesh, manual_axes, check_vma, new_in_specs[i], v)\n\n  with (jshmap._extend_axis_env(mesh, manual_axes),\n        mesh_lib.use_abstract_mesh(jshmap._as_manual_mesh(mesh, manual_axes)),\n        config._check_vma(check_vma)):\n    # jaxpr to checked_jaxpr","sourceCodeStart":960,"sourceCodeEnd":996,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/checkify.py#L960-L996","documentation":"Checkify's shard_map rule requires the mesh keyword, but the transformation was invoked without one. The error values must be re-sharded across the mesh, so the mesh is mandatory.","triggerScenarios":"A shard_map primitive processed under checkify where kwargs lack 'mesh' — typically from custom shard_map usage or a stale/3rd-party path that doesn't thread mesh through.","commonSituations":"Using checkify together with shard_map/jax.shard_map in multi-host or SPMD code; version mismatches where the mesh kwarg isn't propagated by an older helper.","solutions":["Pass mesh explicitly to shard_map / ensure you use jax.shard_map with a Mesh via jax.lax.with_sharding_constraint-style APIs","Upgrade JAX so internal plumbing passes mesh to the checkify rule","Move checkify inside/outside shard_map so both don't need to compose"],"exampleFix":"# before\nout = shard_map(f, mesh=None, in_specs=..., out_specs=...)(x)  # under checkify\n# after\nmesh = jax.sharding.Mesh(jax.devices(), 'd')\nout = jax.shard_map(f, mesh, in_specs=P('d'), out_specs=P('d'))(x)","handlingStrategy":"validation","validationCode":"def safe_shard_map(f, mesh, **kw):\n    if mesh is None:\n        raise ValueError('mesh is required when composing with checkify')\n    return jax.shard_map(f, mesh, **kw)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always construct shard_map from jax.shard_map with an explicit Mesh","Keep JAX version in sync across deps so mesh kwarg propagates"],"tags":["jax","checkify","shard-map","mesh","sharding"],"backgroundTag":"missing-required-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}