{"record":{"id":"61f0ff2be0fb3f8a","repo":"jax-ml/jax","slug":"unreduced-for-max-is-not-yet-supported","errorCode":null,"errorMessage":"unreduced for max is not yet supported.","messagePattern":"unreduced for max is not yet supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/ops/scatter.py","lineNumber":391,"sourceCode":"    segment maximums.\n\n  Examples:\n    Simple 1D segment max:\n\n    >>> data = jnp.arange(6)\n    >>> segment_ids = jnp.array([0, 0, 1, 1, 2, 2])\n    >>> segment_max(data, segment_ids)\n    Array([1, 3, 5], dtype=int32)\n\n    Using JIT requires static `num_segments`:\n\n    >>> from jax import jit\n    >>> jit(segment_max, static_argnums=2)(data, segment_ids, 3)\n    Array([1, 3, 5], dtype=int32)\n  \"\"\"\n  out_sharding = canonicalize_sharding(out_sharding, 'segment_max')\n  if out_sharding is not None and out_sharding.spec.unreduced:\n    raise NotImplementedError('unreduced for max is not yet supported.')\n  return _segment_update(\n      \"segment_max\", data, segment_ids, slicing.scatter_max, num_segments,\n      indices_are_sorted, unique_indices, bucket_size, reductions.max,\n      mode=mode, out_sharding=out_sharding)\n\n\ndef segment_min(data: ArrayLike,\n                segment_ids: ArrayLike,\n                num_segments: int | None = None,\n                indices_are_sorted: bool = False,\n                unique_indices: bool = False,\n                bucket_size: int | None = None,\n                mode: slicing.GatherScatterMode | str | None = None,\n                out_sharding: NamedSharding | P | None = None) -> Array:\n  \"\"\"Computes the minimum within segments of an array.\n\n  Similar to TensorFlow's `segment_min\n  <https://www.tensorflow.org/api_docs/python/tf/math/segment_min>`_","sourceCodeStart":373,"sourceCodeEnd":409,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/ops/scatter.py#L373-L409","documentation":"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.","triggerScenarios":"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).","commonSituations":"Adapting distributed segment_sum code to segment_max; assuming uniform sharding support across the segment op family.","solutions":["Remove out_sharding / pass None for segment_max","Compute segment_max on replicated data, or follow with an explicit reshard","Request the feature upstream if needed"],"exampleFix":"# before\nout = jax.ops.segment_max(data, ids, out_sharding=unreduced_sharding)\n# after\nout = jax.ops.segment_max(data, ids)","handlingStrategy":"validation","validationCode":"if sharding is not None and getattr(sharding.spec, 'unreduced', False):\n    sharding = None  # segment_max has no unreduced lowering\nout = jax.ops.segment_max(data, ids, out_sharding=sharding)","typeGuard":"def supports_unreduced(op_name): return op_name == 'segment_sum'","tryCatchPattern":"catch NotImplementedError, fall back to out_sharding=None","preventionTips":["Document per-op sharding support in your distributed layer","Add unit tests for each segment op x sharding combination you use"],"tags":["jax","sharding","notimplementederror","distributed"],"backgroundTag":"unsupported-feature-flag","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}