jax-ml/jax · error · NotImplementedError

unreduced rule for {name} is not implemented. Please file an

Error message

unreduced rule for {name} is not implemented. Please file an issue at https://github.com/jax-ml/jax/issues

What it means

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.

Source

Thrown at jax/_src/lax/lax.py:4276

  if not any(dtypes.issubdtype(aval.dtype, t) for t in accepted_dtypes):
    msg = '{} does not accept dtype {}. Accepted dtypes are subtypes of {}.'
    typename = dtype_to_string(aval.dtype)
    accepted_typenames = (t.__name__ for t in accepted_dtypes)
    raise TypeError(msg.format(name, typename, ', '.join(accepted_typenames)))
  if (not supports_narrow_ints) and aval.dtype in [dtypes.uint2, dtypes.int2, dtypes.uint4, dtypes.int4]:
    raise TypeError(f'{name} does not accept dtype {dtype_to_string(aval.dtype)}.'
                    ' Support for narrow-width integers is platform-dependent'
                    ' and limited to a few specific operations, e.g. basic'
                    ' arithmetic and type casting.')
  return result_dtype(aval.dtype, **kwargs)

def default_unop_reduced_rule(aval):
  return getr(aval)

def unop_ur_rule(name, aval, **kwargs):
  reduced = default_unop_reduced_rule(aval)
  if any(getu(aval)):
    raise NotImplementedError(
        f'unreduced rule for {name} is not implemented. Please'
        ' file an issue at https://github.com/jax-ml/jax/issues')
  return frozenset(), reduced, None

def unop(result_dtype, accepted_dtypes, name, supports_narrow_ints=True):
  dtype_rule = partial(unop_dtype_rule, result_dtype, accepted_dtypes, name,
                       supports_narrow_ints=supports_narrow_ints)
  prim = standard_primitive(_attrgetter('shape'), dtype_rule, name,
                            sharding_rule=_attrgetter('sharding'),
                            vma_rule=lambda x, **kwargs: x.mat.varying,
                            ur_rule=partial(unop_ur_rule, name))
  batching.defvectorized(prim)
  return prim

standard_unop = partial(unop, _identity)

_attrgetter = lambda name: lambda x, **kwargs: getattr(x, name)

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
@jax.jit(in_shardings=(P('x', None),))
def f(a): return lax.some_op(a)  # hits unimplemented ur rule
// after
def f(a): return lax.some_op(jax.lax.with_sharding_constraint(a, P()))

Defensive patterns

Strategy: fallback

Try / catch

try:
    out = sharded_jit_fn(x)
except NotImplementedError as e:
    if 'unreduced rule' in str(e):
        out = unsharded_fn(jax.device_get(x))  # fallback path
    else:
        raise

Prevention

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

When it happens

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

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

Related errors

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