jax-ml/jax · error · NotImplementedError

unreduced for prod is not yet supported.

Error message

unreduced for prod is not yet supported.

What it means

jax.ops.segment_prod does not support unreduced output sharding. When the out_sharding argument resolves to a sharding whose spec has unreduced=True, JAX raises NotImplementedError because the product reduction has no unreduced lowering implemented (unlike segment_sum).

Source

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

    segment products.

  Examples:
    Simple 1D segment product:

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

    Using JIT requires static `num_segments`:

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


def segment_max(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 maximum within segments of an array.

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the out_sharding argument (or pass None) for segment_prod
2. If distributed output is required, use segment_sum where unreduced is supported, or restructure with manual psum-style reductions
3. File/track the feature request upstream for unreduced segment_prod

Example fix

# before
out = jax.ops.segment_prod(data, ids, out_sharding=unreduced_sharding)
# after
out = jax.ops.segment_prod(data, ids)  # materialize fully, reshard afterwards

sh = out_sharding
if sh is not None and getattr(sh.spec, 'unreduced', False):
    raise ValueError('segment_prod has no unreduced lowering; use None sharding')
out = jax.ops.segment_prod(data, ids, out_sharding=None)

def is_reduced_sharding(s):
    return s is None or not getattr(s.spec, 'unreduced', False)

catch NotImplementedError and retry without out_sharding

When it happens

Trigger: Calling jax.ops.segment_prod(data, segment_ids, out_sharding=sharding) where sharding (e.g. from jax.sharding.NamedSharding or GSPMDSharding) has .spec.unreduced true, typically in multi-host TPU pipelines.

Common situations: Copy-pasting a segment_sum call with unreduced out_sharding to segment_prod; migrating SPMD distributed code expecting all segment ops to support unreduced output.

Related errors

• unreduced for max is not yet supported.
• 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

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