jax-ml/jax · error · NotImplementedError

unreduced for max is not yet supported.

Error message

unreduced for max is not yet supported.

What it means

jax.ops.segment_max does not support unreduced output sharding. If out_sharding canonicalizes to a spec with unreduced=True, JAX raises NotImplementedError since only segment_sum has an unreduced lowering; max reductions across hosts would require cross-device communication that is not implemented.

Source

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

    segment maximums.

  Examples:
    Simple 1D segment max:

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

    Using JIT requires static `num_segments`:

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


def segment_min(data: ArrayLike,
                segment_ids: ArrayLike,
                num_segments: int | None = None,
                indices_are_sorted: bool = False,
                unique_indices: bool = False,
                bucket_size: int | None = None,
                mode: slicing.GatherScatterMode | str | None = None,
                out_sharding: NamedSharding | P | None = None) -> Array:
  """Computes the minimum within segments of an array.

  Similar to TensorFlow's `segment_min
  <https://www.tensorflow.org/api_docs/python/tf/math/segment_min>`_

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove out_sharding / pass None for segment_max
2. Compute segment_max on replicated data, or follow with an explicit reshard
3. Request the feature upstream if needed

Example fix

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

if sharding is not None and getattr(sharding.spec, 'unreduced', False):
    sharding = None  # segment_max has no unreduced lowering
out = jax.ops.segment_max(data, ids, out_sharding=sharding)

def supports_unreduced(op_name): return op_name == 'segment_sum'

catch NotImplementedError, fall back to out_sharding=None

When it happens

Trigger: Calling jax.ops.segment_max(data, segment_ids, ..., out_sharding=s) where s.spec.unreduced is True (GSPMD-style sharding annotations on TPU/multi-device).

Common situations: Adapting distributed segment_sum code to segment_max; assuming uniform sharding support across the segment op family.

Related errors

• unreduced for min 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/61f0ff2be0fb3f8a. Report an issue: GitHub.