{"record":{"id":"cddcc982dc9273b8","repo":"jax-ml/jax","slug":"unreduced-for-prod-is-not-yet-supported","errorCode":null,"errorMessage":"unreduced for prod is not yet supported.","messagePattern":"unreduced for prod is not yet supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/ops/scatter.py","lineNumber":332,"sourceCode":"    segment products.\n\n  Examples:\n    Simple 1D segment product:\n\n    >>> data = jnp.arange(6)\n    >>> segment_ids = jnp.array([0, 0, 1, 1, 2, 2])\n    >>> segment_prod(data, segment_ids)\n    Array([ 0,  6, 20], dtype=int32)\n\n    Using JIT requires static `num_segments`:\n\n    >>> from jax import jit\n    >>> jit(segment_prod, static_argnums=2)(data, segment_ids, 3)\n    Array([ 0,  6, 20], dtype=int32)\n  \"\"\"\n  out_sharding = canonicalize_sharding(out_sharding, 'segment_prod')\n  if out_sharding is not None and out_sharding.spec.unreduced:\n    raise NotImplementedError('unreduced for prod is not yet supported.')\n  return _segment_update(\n      \"segment_prod\", data, segment_ids, slicing.scatter_mul, num_segments,\n      indices_are_sorted, unique_indices, bucket_size, reductions.prod,\n      mode=mode, out_sharding=out_sharding)\n\n\ndef segment_max(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 maximum within segments of an array.\n\n  Similar to TensorFlow's `segment_max\n  <https://www.tensorflow.org/api_docs/python/tf/math/segment_max>`_","sourceCodeStart":314,"sourceCodeEnd":350,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/ops/scatter.py#L314-L350","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["Drop the out_sharding argument (or pass None) for segment_prod","If distributed output is required, use segment_sum where unreduced is supported, or restructure with manual psum-style reductions","File/track the feature request upstream for unreduced segment_prod"],"exampleFix":"# before\nout = jax.ops.segment_prod(data, ids, out_sharding=unreduced_sharding)\n# after\nout = jax.ops.segment_prod(data, ids)  # materialize fully, reshard afterwards","handlingStrategy":"validation","validationCode":"sh = out_sharding\nif sh is not None and getattr(sh.spec, 'unreduced', False):\n    raise ValueError('segment_prod has no unreduced lowering; use None sharding')\nout = jax.ops.segment_prod(data, ids, out_sharding=None)","typeGuard":"def is_reduced_sharding(s):\n    return s is None or not getattr(s.spec, 'unreduced', False)","tryCatchPattern":"catch NotImplementedError and retry without out_sharding","preventionTips":["Only use unreduced shardings with segment_sum","Centralize sharding choices in one config so unsupported combos are caught early"],"tags":["jax","pallas-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"}