jax-ml/jax · error · NotImplementedError

unreduced for min is not yet supported.

Error message

unreduced for min is not yet supported.

What it means

jax.ops.segment_min does not support unreduced output sharding. Passing an out_sharding whose spec sets unreduced=True triggers NotImplementedError because only segment_sum implements the unreduced lowering.

Source

Thrown at jax/_src/ops/scatter.py:450

    segment minimums.

  Examples:
    Simple 1D segment min:

    >>> data = jnp.arange(6)
    >>> segment_ids = jnp.array([0, 0, 1, 1, 2, 2])
    >>> segment_min(data, segment_ids)
    Array([0, 2, 4], dtype=int32)

    Using JIT requires static `num_segments`:

    >>> from jax import jit
    >>> jit(segment_min, static_argnums=2)(data, segment_ids, 3)
    Array([0, 2, 4], dtype=int32)
  """
  out_sharding = canonicalize_sharding(out_sharding, 'segment_min')
  if out_sharding is not None and out_sharding.spec.unreduced:
    raise NotImplementedError('unreduced for min is not yet supported.')
  return _segment_update(
      "segment_min", data, segment_ids, slicing.scatter_min, num_segments,
      indices_are_sorted, unique_indices, bucket_size, reductions.min,
      mode=mode, out_sharding=out_sharding)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass out_sharding=None for segment_min
2. Compute locally then redistribute the result with jax.device_put
3. Track upstream support for unreduced min

Example fix

# before
out = jax.ops.segment_min(data, ids, out_sharding=unreduced_sharding)
# after
out = jax.ops.segment_min(data, ids)

if sharding is not None and getattr(sharding.spec, 'unreduced', False):
    sharding = None
out = jax.ops.segment_min(data, ids, out_sharding=sharding)

catch NotImplementedError and recompute without sharding

When it happens

Trigger: Calling jax.ops.segment_min(data, segment_ids, out_sharding=sharding) with an unreduced sharding spec in a multi-device/TPU program.

Common situations: Reusing sharding configs written for segment_sum with segment_min; distributed training pipelines using GSPMD annotations.

Related errors

• unreduced for max is not yet supported.
• Mapped away dimension of inputs passed to vmap should be sha
• The 'sharding' attribute is not available on {self._error_re
• The is_fully_addressable property was called on {self._error
• The error code state and the predicate must be on the same m

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