jax-ml/jax · error · NotImplementedError

Eager shard_map + unreduced/reduced + partial manual is not

Error message

Eager shard_map + unreduced/reduced + partial manual is not implemented. Please wrap your shard_map in `jax.jit`.

What it means

In eager (non-jit) execution with partial-manual axes, shard_map has no implementation for inputs whose specs carry unreduced or reduced fields; the machinery to reconcile them outside of tracing is not built. It raises NotImplementedError and asks you to use jax.jit.

Source

Thrown at jax/_src/shard_map.py:1425

      out_mat = (out_mat if isinstance(out_mat, (list, tuple))
                else [out_mat] * len(out_vals))
      return map(partial(ShardMapTracer, self), out_mat, out_vals)
    return ShardMapTracer(self, out_mat, out_vals)

  def process_shard_map(self, prim, fun, args, mesh, in_specs,
                        check_vma, newly_manual_axes, debug_info):
    # Check consistency between outer and inner shmaps on explicitly passed
    # mesh and check_vma.
    if isinstance(mesh, Mesh):
      if mesh != self.mesh: raise Exception
    del mesh
    if check_vma != self.check:  # TODO(mattjj): add check in jit path
      raise Exception
    del check_vma

    in_vals, in_mats = unzip2(map(self.to_val_mat_pair, args))
    if any(m.unreduced or m.reduced for m in in_mats):
      raise NotImplementedError(
          "Eager shard_map + unreduced/reduced + partial manual is not"
          " implemented. Please wrap your shard_map in `jax.jit`.")
    trace = ShardMapTrace(self.mesh, newly_manual_axes | self.manual_axes, self.check)
    in_vals_ = [_unmatch_spec2(self.mesh, self.manual_axes, spec, x)
                for x, spec in zip(in_vals, in_specs)]
    # TODO(yashkatariya): Handle unreduced/reduced correctly.
    in_mats_ = [core.ManualAxisType(varying=mat.varying | _spec_to_vma(s))
                for mat, s in zip(in_mats, in_specs)]
    in_tracers = map(partial(ShardMapTracer, trace), in_mats_, in_vals_)
    inner_mesh = _as_manual_mesh(self.mesh, newly_manual_axes | self.manual_axes)
    with (core.set_current_trace(trace), _extend_axis_env(self.mesh, newly_manual_axes),
          use_abstract_mesh(inner_mesh)):
      ans_aux = fun(*in_tracers)
      ans, out_specs = ans_aux.unpack_aux()
      out_vals_, out_mats_ = ans.map(trace.to_val_mat_pair).unzip2()
    out_vals = out_vals_.map2(
        out_specs, lambda x, spec: _match_spec2(self.mesh, self.manual_axes, spec, x))
    # TODO(yashkatariya): Handle unreduced/reduced correctly.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Wrap the shard_map call in jax.jit (the supported path)
Remove unreduced/reduced from input specs if eager execution is required
Drop partial manual_axes (use full manual) if semantics allow

Example fix

# before
f = shard_map(body, mesh, manual_axes=('m',), in_specs=P(reduced=('r',)))
y = f(x)
# after
f = jax.jit(shard_map(body, mesh, manual_axes=('m',), in_specs=P(reduced=('r',))))
y = f(x)

Defensive patterns

Strategy: validation

Validate before calling

def jit_required(specs, manual_axes, mesh):
    uses_exotic = any(getattr(p, 'unreduced', None) or getattr(p, 'reduced', None) for p in jax.tree.leaves(specs))
    return uses_exotic and set(manual_axes) != set(mesh.axis_names)

Try / catch

try: f(x) except NotImplementedError as e: if 'unreduced/reduced' in str(e): f = jax.jit(f); f(x); else: raise

Prevention

Always jit shard_map functions using unreduced/reduced specs
Test eager paths separately without partial-manual specs

When it happens

Trigger: Calling an un-jitted shard_map with manual_axes a strict subset of mesh axes, where any input spec/aval has non-empty unreduced or reduced.

Common situations: Interactive experimentation / notebook calls without jit while prototyping SP-style partial-reduce code.

Related errors

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