{"record":{"id":"8e61211f1c4921ea","repo":"jax-ml/jax","slug":"unreduced-rule-for-name-is-not-implemented-plea","errorCode":null,"errorMessage":"unreduced rule for {name} is not implemented. Please file an issue at https://github.com/jax-ml/jax/issues","messagePattern":"unreduced rule for (.+?) is not implemented\\. Please file an issue at https://github\\.com/jax-ml/jax/issues","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":4276,"sourceCode":"  if not any(dtypes.issubdtype(aval.dtype, t) for t in accepted_dtypes):\n    msg = '{} does not accept dtype {}. Accepted dtypes are subtypes of {}.'\n    typename = dtype_to_string(aval.dtype)\n    accepted_typenames = (t.__name__ for t in accepted_dtypes)\n    raise TypeError(msg.format(name, typename, ', '.join(accepted_typenames)))\n  if (not supports_narrow_ints) and aval.dtype in [dtypes.uint2, dtypes.int2, dtypes.uint4, dtypes.int4]:\n    raise TypeError(f'{name} does not accept dtype {dtype_to_string(aval.dtype)}.'\n                    ' Support for narrow-width integers is platform-dependent'\n                    ' and limited to a few specific operations, e.g. basic'\n                    ' arithmetic and type casting.')\n  return result_dtype(aval.dtype, **kwargs)\n\ndef default_unop_reduced_rule(aval):\n  return getr(aval)\n\ndef unop_ur_rule(name, aval, **kwargs):\n  reduced = default_unop_reduced_rule(aval)\n  if any(getu(aval)):\n    raise NotImplementedError(\n        f'unreduced rule for {name} is not implemented. Please'\n        ' file an issue at https://github.com/jax-ml/jax/issues')\n  return frozenset(), reduced, None\n\ndef unop(result_dtype, accepted_dtypes, name, supports_narrow_ints=True):\n  dtype_rule = partial(unop_dtype_rule, result_dtype, accepted_dtypes, name,\n                       supports_narrow_ints=supports_narrow_ints)\n  prim = standard_primitive(_attrgetter('shape'), dtype_rule, name,\n                            sharding_rule=_attrgetter('sharding'),\n                            vma_rule=lambda x, **kwargs: x.mat.varying,\n                            ur_rule=partial(unop_ur_rule, name))\n  batching.defvectorized(prim)\n  return prim\n\nstandard_unop = partial(unop, _identity)\n\n_attrgetter = lambda name: lambda x, **kwargs: getattr(x, name)\n","sourceCodeStart":4258,"sourceCodeEnd":4294,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L4258-L4294","documentation":"JAX's sharding/partitioning propagation defines 'unreduced' rules for when an operand is unreduced along some axes. The default unop unreduced rule does not support operands that are unreduced, so it raises NotImplementedError and asks for a GitHub issue.","triggerScenarios":"Automatic sharding propagation (e.g. GSPMD / NamedSharded jits) where a unop's operand has unreduced axes and no custom ur_rule was registered for that primitive.","commonSituations":"Emerging sharding APIs: combining named shardings with ops whose sharding rules are incomplete; upgrading JAX versions where a newly sharded pipeline hits an unimplemented path.","solutions":["Report it: file an issue at https://github.com/jax-ml/jax/issues with a minimal reproducer","Restructure so the operand is fully reduced before the op (reshard/gather to replicated)","Work around by materializing the array (jit with sharding disabled for that section) or using an equivalent op with implemented sharding rules"],"exampleFix":"// before\n@jax.jit(in_shardings=(P('x', None),))\ndef f(a): return lax.some_op(a)  # hits unimplemented ur rule\n// after\ndef f(a): return lax.some_op(jax.lax.with_sharding_constraint(a, P()))","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    out = sharded_jit_fn(x)\nexcept NotImplementedError as e:\n    if 'unreduced rule' in str(e):\n        out = unsharded_fn(jax.device_get(x))  # fallback path\n    else:\n        raise","preventionTips":["Test sharded pipelines on small configs before scaling","Pin JAX version once a sharding path is validated","Constrain shardings explicitly instead of relying on propagation"],"tags":["jax","sharding","gspmd","not-implemented","unreduced"],"backgroundTag":"unimplemented-sharding-rule","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}