{"record":{"id":"4127332175d57e9e","repo":"jax-ml/jax","slug":"eager-shard-map-unreduced-reduced-partial-manu","errorCode":null,"errorMessage":"Eager shard_map + unreduced/reduced + partial manual is not implemented. Please wrap your shard_map in `jax.jit`.","messagePattern":"Eager shard_map \\+ unreduced/reduced \\+ partial manual is not implemented\\. Please wrap your shard_map in `jax\\.jit`\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1425,"sourceCode":"      out_mat = (out_mat if isinstance(out_mat, (list, tuple))\n                else [out_mat] * len(out_vals))\n      return map(partial(ShardMapTracer, self), out_mat, out_vals)\n    return ShardMapTracer(self, out_mat, out_vals)\n\n  def process_shard_map(self, prim, fun, args, mesh, in_specs,\n                        check_vma, newly_manual_axes, debug_info):\n    # Check consistency between outer and inner shmaps on explicitly passed\n    # mesh and check_vma.\n    if isinstance(mesh, Mesh):\n      if mesh != self.mesh: raise Exception\n    del mesh\n    if check_vma != self.check:  # TODO(mattjj): add check in jit path\n      raise Exception\n    del check_vma\n\n    in_vals, in_mats = unzip2(map(self.to_val_mat_pair, args))\n    if any(m.unreduced or m.reduced for m in in_mats):\n      raise NotImplementedError(\n          \"Eager shard_map + unreduced/reduced + partial manual is not\"\n          \" implemented. Please wrap your shard_map in `jax.jit`.\")\n    trace = ShardMapTrace(self.mesh, newly_manual_axes | self.manual_axes, self.check)\n    in_vals_ = [_unmatch_spec2(self.mesh, self.manual_axes, spec, x)\n                for x, spec in zip(in_vals, in_specs)]\n    # TODO(yashkatariya): Handle unreduced/reduced correctly.\n    in_mats_ = [core.ManualAxisType(varying=mat.varying | _spec_to_vma(s))\n                for mat, s in zip(in_mats, in_specs)]\n    in_tracers = map(partial(ShardMapTracer, trace), in_mats_, in_vals_)\n    inner_mesh = _as_manual_mesh(self.mesh, newly_manual_axes | self.manual_axes)\n    with (core.set_current_trace(trace), _extend_axis_env(self.mesh, newly_manual_axes),\n          use_abstract_mesh(inner_mesh)):\n      ans_aux = fun(*in_tracers)\n      ans, out_specs = ans_aux.unpack_aux()\n      out_vals_, out_mats_ = ans.map(trace.to_val_mat_pair).unzip2()\n    out_vals = out_vals_.map2(\n        out_specs, lambda x, spec: _match_spec2(self.mesh, self.manual_axes, spec, x))\n    # TODO(yashkatariya): Handle unreduced/reduced correctly.","sourceCodeStart":1407,"sourceCodeEnd":1443,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1407-L1443","documentation":"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.","triggerScenarios":"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.","commonSituations":"Interactive experimentation / notebook calls without jit while prototyping SP-style partial-reduce code.","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"],"exampleFix":"# before\nf = shard_map(body, mesh, manual_axes=('m',), in_specs=P(reduced=('r',)))\ny = f(x)\n# after\nf = jax.jit(shard_map(body, mesh, manual_axes=('m',), in_specs=P(reduced=('r',))))\ny = f(x)","handlingStrategy":"validation","validationCode":"def jit_required(specs, manual_axes, mesh):\n    uses_exotic = any(getattr(p, 'unreduced', None) or getattr(p, 'reduced', None) for p in jax.tree.leaves(specs))\n    return uses_exotic and set(manual_axes) != set(mesh.axis_names)","typeGuard":null,"tryCatchPattern":"try: f(x) except NotImplementedError as e: if 'unreduced/reduced' in str(e): f = jax.jit(f); f(x); else: raise","preventionTips":["Always jit shard_map functions using unreduced/reduced specs","Test eager paths separately without partial-manual specs"],"tags":["shard-map","eager","not-implemented","unreduced","jax"],"backgroundTag":"eager-mode-unsupported-op","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}