jax-ml/jax · error · ValueError

Mesh must be provided for shard_map with checkify.

Error message

Mesh must be provided for shard_map with checkify.

What it means

checkify's shard_map error-check requires knowledge of the mesh over which the sharded computation runs. The mesh parameter was not forwarded in the call kwargs, so checkify cannot construct error-value shardings and refuses to proceed.

Solutions

  1. Pass mesh explicitly to shard_map / ensure you use jax.shard_map with a Mesh via jax.lax.with_sharding_constraint-style APIs
  2. Upgrade JAX so internal plumbing passes mesh to the checkify rule
  3. Move checkify inside/outside shard_map so both don't need to compose

Example fix

# before
out = shard_map(f, mesh=None, in_specs=..., out_specs=...)(x)  # under checkify
# after
mesh = jax.sharding.Mesh(jax.devices(), 'd')
out = jax.shard_map(f, mesh, in_specs=P('d'), out_specs=P('d'))(x)
Defensive patterns

Strategy: validation

Validate before calling

def safe_shard_map(f, mesh, **kw):
    if mesh is None:
        raise ValueError('mesh is required when composing with checkify')
    return jax.shard_map(f, mesh, **kw)

Prevention

When it happens

Trigger: 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.

Common situations: 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.

Understand the failure class

Background: "missing required argument" and "the following required arguments were not provided": what required-argument errors mean and how to fix them — this error's family across 20 libraries.

Related errors

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

Appendix: source

Thrown at jax/_src/checkify.py:978

                                 *vals_in, _prim=_prim)
  err_vals, err_tree = jtu.tree_flatten(error)
  new_vals_in = [*err_vals, *vals_in]
  in_avals = tuple(map(core.typeof, new_vals_in))
  checked_jaxpr_, out_tree, _ = jaxpr_to_checkify_jaxpr(
      _prim.jaxpr, enabled_errors, err_tree, *in_avals)
  checked_jaxpr, consts = pe.separate_consts(checked_jaxpr_)
  new_prim = ad_checkpoint.RematTraced(checked_jaxpr, _prim.policy)
  err_and_out = new_prim(*consts, *new_vals_in)
  return tree_unflatten(out_tree, err_and_out)
error_checks[call_hi_primitive_p] = call_hi_primitive_error_check


def shard_map_error_check(
    error: Error, enabled_errors, *vals_in,
    jaxpr: core.Jaxpr, in_specs, out_specs, **kwargs
):
  if (mesh := kwargs.get('mesh')) is None:
    raise ValueError('Mesh must be provided for shard_map with checkify.')

  err_vals, err_tree = jtu.tree_flatten(error)
  num_error_vals = len(err_vals)
  # Replicated sharding for in errors.
  new_in_specs = (*([P()] * num_error_vals), *in_specs)
  new_vals_in = [*err_vals, *vals_in]
  in_avals = list(map(core.typeof, new_vals_in))
  manual_axes = kwargs.get('newly_manual_axes')
  check_vma = kwargs.get('check_vma')
  for i, v in enumerate(in_avals):
    if not (sharder := core.shard_aval_handlers.get(type(v))):
      raise ValueError(f'Unsupported aval type: {type(v)}')
    in_avals[i] = sharder(mesh, manual_axes, check_vma, new_in_specs[i], v)

  with (jshmap._extend_axis_env(mesh, manual_axes),
        mesh_lib.use_abstract_mesh(jshmap._as_manual_mesh(mesh, manual_axes)),
        config._check_vma(check_vma)):
    # jaxpr to checked_jaxpr

View on GitHub (pinned to 1e1c6a8fc0)