[xpu-alignment] [26-10-21] XPU segment_reduce max/min backward gives a negative tied gradient to every tied element instead of splitting it

Summary

The backward of torch.segment_reduce / torch._segment_reduce with reduce="max" or reduce="min" on XPU splits the incoming gradient between the elements tied at the segment extremum only when that gradient is positive. For a negative gradient every tied element receives the full gradient, so the segment's input gradients sum to k * grad_output instead of grad_output. XPU amax / amin backward splits evenly for both signs.

The error is silent, sign-dependent and therefore biased (it does not average out over training steps). gradcheck does not see it, because random inputs have no ties, and the existing tests only call .sum().backward(), i.e. use a positive incoming gradient.

Upstream source

segment_reduce max/min backward does not split negative gradients between tied elements pytorch/pytorch#201633 (open, filed 2026-10-21) reports the same defect for the CPU and CUDA kernels.
Found by an XPU alignment scan of the window [2026-10-21T00:00:00Z, 2026-10-22T00:00:00Z) over pytorch/pytorch, frozen default-branch head 37f56e48ce431907080585a595c437f9edb23411.

Observed XPU behavior

Validated on XPU with torch 2.16.0.dev20261021+xpu:

case reduce=max grad_out=+1.0 tied=2 grad=[0.5, 0.5, 0.0] grad_sum=+1.000000 amax_amin_grad=[0.5, 0.5, 0.0] amax_amin_sum=+1.000000 matches_reference=True
case reduce=max grad_out=-1.0 tied=2 grad=[-1.0, -1.0, 0.0] grad_sum=-2.000000 amax_amin_grad=[-0.5, -0.5, -0.0] amax_amin_sum=-1.000000 matches_reference=False
case reduce=min grad_out=+1.0 tied=2 grad=[0.5, 0.5, 0.0] grad_sum=+1.000000 amax_amin_grad=[0.5, 0.5, 0.0] amax_amin_sum=+1.000000 matches_reference=True
case reduce=min grad_out=-1.0 tied=2 grad=[-1.0, -1.0, 0.0] grad_sum=-2.000000 amax_amin_grad=[-0.5, -0.5, -0.0] amax_amin_sum=-1.000000 matches_reference=False

Expected for grad_out=-1: [-0.5, -0.5, 0.0], as for amax / amin. The positive-gradient rows already split evenly in the same kernel, so the kernel is inconsistent with itself as well as with amax / amin.

Target path and root cause

segment_reduce_backward_kernel in src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp (line 290). In the MAX/MIN branch (line 338) the first loop selects the tied elements with isnan(values_data[data_index]) || values_data[data_index] == output_data[output_index] and counts them, and the second loop then divides by that count:

// src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp:359
if (grad_input_data[data_index] > 0) {
  grad_input_data[data_index] = grad_input_data[data_index] / counter;
}

The division loop gates on the sign of the value it just wrote instead of re-using the tie predicate, so only positive gradients are divided. The fix is to select the tied elements in the division loop with the same predicate used in the assignment loop.

This kernel backs both _segment_reduce_lengths_backward_kernel (line 554) and _segment_reduce_offsets_backward_kernel (line 573), so lengths- and offsets-based calls are both affected.

Target-path evidence: aten::_segment_reduce_backward has an XPU kernel registered (dispatch_segment_reduce_backward_xpu True), and both the forward output and the accumulated .grad are xpu tensors. The amax / amin reference was computed from the same XPU inputs on the same device, so the discrepancy is between two XPU paths and is not a host/device comparison artifact.

Reproducer

import torch

for reduce, vals in (("max", [2.0, 2.0, 1.0]), ("min", [1.0, 1.0, 2.0])):
    for g in (1.0, -1.0):
        data = torch.tensor(vals, device="xpu", requires_grad=True)
        out = torch.segment_reduce(data, reduce, lengths=torch.tensor([3], device="xpu"))
        out.backward(torch.tensor([g], device="xpu"))

        ref = torch.tensor(vals, device="xpu", requires_grad=True)
        (ref.amax() if reduce == "max" else ref.amin()).backward(torch.tensor(g, device="xpu"))

        print(f"{reduce} grad_out={g:+.0f}  segment_reduce: {data.grad.tolist()} "
              f"(sum {data.grad.sum().item():+.1f})   "
              f"{'amax' if reduce == 'max' else 'amin'}: {ref.grad.tolist()}")

Environment

torch 2.16.0.dev20261021+xpu
python 3.13.13
device Intel(R) Arc(TM) Pro B60 Graphics

Ownership

The XPU kernel carries its own copy of the gate. The same grad_input_data[data_index] > 0 gate is present upstream in aten/src/ATen/native/SegmentReduce.cpp:271 and aten/src/ATen/native/cuda/SegmentReduce.cu:222 at the scanned head, so an upstream fix to the CPU and CUDA kernels will not change src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp; the XPU change is independent work in this repository.

Current stage of the related upstream fix: pytorch/pytorch#201633 is open and the reporter states they have a fix and will open a PR; no fix PR is linked to the issue as of this scan, so the upstream change is proposed, not landed. The XPU fix should adopt the same tie predicate so XPU matches the CPU/CUDA convention once that PR lands.

Partially related, not covering this finding: pytorch/pytorch#193476 (open) bounds-checks offsets in the CUDA segment_reduce_backward_kernel to fix an out-of-bounds write; it mentions the tie split only in a review comment and does not change the gate. No open or closed intel/torch-xpu-ops tracker was found for the tie split (searched segment_reduce, segment reduce backward, tied gradient, and PRs touching SegmentReduce*), and upstream test/test_segment_reductions.py carries no XPU skip, xfail or tracker reference at the scanned head.

Regression test

No existing test covers this: the upstream cases use a positive incoming gradient. A fix should add a tied-segment case with a negative grad_output for both max and min, asserting agreement with amax / amin backward on the same inputs.

@torchxpubot

@torchxpubot fix

Automated fix session for #5808. Stages are appended below as they complete.

Issue Triage

Field	Value
Issue type	single-bug
Reproduction missing	no
Scope	torch-xpu-ops
Dependencies	(none)
Handling	agent-fixable

Automated by issue-triage.

Reproduce

Reproduced against the nightly wheel torch 2.16.0.dev20261027+xpu on Intel(R) Arc(TM) Pro B60 Graphics (BMG). The reproducer asserts the segment_reduce backward gradient against amax/amin backward on the same XPU inputs, and covers both the lengths= and offsets= entry points.

Test case	Verdict	Observed
`reduce=max grad_out=-1 lengths`	REPRODUCED	`segment_reduce=[-1.0, -1.0, 0.0] (sum -2.0) expected=[-0.5, -0.5, -0.0]`
`reduce=max grad_out=-1 offsets`	REPRODUCED	`segment_reduce=[-1.0, -1.0, 0.0] (sum -2.0) expected=[-0.5, -0.5, -0.0]`
`reduce=min grad_out=-1 lengths`	REPRODUCED	`segment_reduce=[-1.0, -1.0, 0.0] (sum -2.0) expected=[-0.5, -0.5, -0.0]`
`reduce=min grad_out=-1 offsets`	REPRODUCED	`segment_reduce=[-1.0, -1.0, 0.0] (sum -2.0) expected=[-0.5, -0.5, -0.0]`

The four positive-grad_output rows of the same matrix pass, confirming the defect is sign-gated rather than a general tie-splitting failure.

Automated by fix-reproduce.

Root Cause

Analyzed at torch-xpu-ops@73a6552c3, pytorch origin/main@ec1459e58.

segment_reduce_backward_kernel in src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp:290 handles MAX/MIN in two passes over the segment. The first pass (line 341) selects the elements tied at the segment extremum with the predicate sycl::isnan(values_data[data_index]) || values_data[data_index] == output_data[output_index], writes the incoming grad_data[output_index] to each of them, and counts them into counter. The second pass (line 356) is supposed to split that gradient counter ways, but it re-selects the elements with if (grad_input_data[data_index] > 0) (line 359) -- the sign of the value it just wrote -- instead of re-using the tie predicate. A negative incoming gradient therefore fails the gate and is never divided, so every one of the k tied elements keeps the full grad_output and the segment's input gradients sum to k * grad_output. A positive gradient takes the same path and is divided correctly, which is why the defect is sign-gated and silent.

This is not a regression: the gate has been present since the kernel was written, and the two tests that exercise max/min backward cannot see it. test/test_segment_reductions.py drives backward through actual_result.sum().backward() (line 93), i.e. an all-ones -- positive -- incoming gradient, and its gradcheck cases (lines 482-483) use src values that are distinct within each segment, so no tie is ever formed. Both the lengths- and offsets-based entry points route through this one kernel (_segment_reduce_lengths_backward_kernel at line 554, _segment_reduce_offsets_backward_kernel at line 573), and the reproduce stage confirmed all four combinations fail.

The same gate exists upstream in aten/src/ATen/native/SegmentReduce.cpp:271 and aten/src/ATen/native/cuda/SegmentReduce.cu:222, so the XPU kernel inherited the defect rather than diverging from a correct reference. pytorch/pytorch#201633 reports it for CPU and CUDA and is still open with no fix PR linked, so no upstream change will touch the XPU file -- the fix here is independent work in this repository. The correct behavior is defined by amax/amin backward, which splits evenly for both signs, and is reached by selecting the tied elements in the division pass with the same predicate the assignment pass uses. No test anywhere covers a tied segment with a negative incoming gradient, so a regression test is part of the fix.

Verdict: IMPLEMENTING -- a two-line predicate correction in one SYCL kernel, plus a regression test
Target repo: torch-xpu-ops
Domain: xpu-kernel

Automated by fix-root-cause.

Implement Result

Target repo: torch-xpu-ops
Analyzed at: torch-xpu-ops@73a6552c3
What I changed: src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp (+3/-1): the MAX/MIN division pass now selects tied elements with the same predicate the assignment pass uses. test/regressions/test_segment_reduce_backward_ties.py (+71): new XPU regression test.
Why: The division pass gated on the sign of the value it had just written (SegmentReduceKernels.cpp:359) rather than on tie membership, so only positive gradients were divided by the tie count. The predicate it now uses is literally the one from the assignment pass at SegmentReduceKernels.cpp:344-346, which is the condition that defines the tie, so both signs are split. The expected result is defined by amax/amin backward, which the new test uses as its reference on the same XPU inputs. The test is XPU-only because the identical gate is still live upstream at aten/src/ATen/native/SegmentReduce.cpp:271 and aten/src/ATen/native/cuda/SegmentReduce.cu:222, so a device-generic version would fail on CPU.
Skip added: no
Ready for verify: yes

Skip-guard review (allow_skip=false): APPROVE.

Unrelated pre-existing issue noticed, not touched: SegmentReduceKernels.cpp:312-317 has the if (idx >= size) { return; } guard duplicated verbatim.

diff --git a/src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp b/src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp
index 0fe0e1320..e7547c41c 100644
--- a/src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp
+++ b/src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp
@@ -356,7 +356,9 @@ void segment_reduce_backward_kernel(
     for (int64_t j = offset_start; j < offset_end; ++j) {
       int64_t data_index = outer_idx * data_stride_axis * data_size_axis +
           j * data_stride_axis + lane_id;
-      if (grad_input_data[data_index] > 0) {
+      if (sycl::isnan(static_cast<at::opmath_type<scalar_t>>(
+              values_data[data_index])) ||
+          values_data[data_index] == output_data[output_index]) {
         grad_input_data[data_index] = grad_input_data[data_index] / counter;
       }
     }
diff --git a/test/regressions/test_segment_reduce_backward_ties.py b/test/regressions/test_segment_reduce_backward_ties.py
new file mode 100644
index 000000000..c841510f1
--- /dev/null
+++ b/test/regressions/test_segment_reduce_backward_ties.py
@@ -0,0 +1,71 @@
+# Copyright 2020-2026 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Owner(s): ["module: intel"]
+"""
+Regression test for segment_reduce max/min backward on tied segments (XPU).
+
+The backward kernel assigns the incoming gradient to every element tied at the
+segment extremum and then divides by the tie count. The division used to be
+gated on the written value being positive, so a negative gradient was handed in
+full to each of the k tied elements and the segment's input gradients summed to
+k * grad_output instead of grad_output (intel/torch-xpu-ops#5808).
+
+amax/amin backward splits evenly for both signs and is the reference here. The
+test stays XPU-only: the same gate is still present in the upstream CPU and
+CUDA kernels (pytorch/pytorch#201633), so a device-generic version would fail
+on CPU.
+"""
+
+import torch
+from torch.testing._internal.common_utils import (
+    instantiate_parametrized_tests,
+    parametrize,
+    run_tests,
+    TestCase,
+)
+
+# Segment 0 has three tied elements, segment 1 none, segment 2 two.
+_VALUES = [3.0, 3.0, 3.0, 1.0, 2.0, 5.0, 5.0]
+_LENGTHS = [3, 2, 2]
+
+
+def _amax_amin_reference_grad(reduce, grad_output):
+    """Per-segment amax/amin backward over the same XPU inputs."""
+    data = torch.tensor(_VALUES, device="xpu", requires_grad=True)
+    outputs = []
+    start = 0
+    for length in _LENGTHS:
+        segment = data[start : start + length]
+        outputs.append(segment.amax() if reduce == "max" else segment.amin())
+        start += length
+    torch.stack(outputs).backward(grad_output)
+    return data.grad
+
+
+@instantiate_parametrized_tests
+class TestSegmentReduceBackwardTies(TestCase):
+    @parametrize("reduce", ["max", "min"])
+    @parametrize("mode", ["lengths", "offsets"])
+    @parametrize("grad_sign", [1.0, -1.0])
+    def test_tied_gradient_is_split_for_both_signs(self, reduce, mode, grad_sign):
+        grad_output = grad_sign * torch.tensor([1.0, 2.0, 3.0], device="xpu")
+
+        data = torch.tensor(_VALUES, device="xpu", requires_grad=True)
+        if mode == "lengths":
+            segments = {"lengths": torch.tensor(_LENGTHS, device="xpu")}
+        else:
+            offsets = torch.tensor(_LENGTHS, device="xpu").cumsum(0)
+            segments = {"offsets": torch.cat([offsets.new_zeros(1), offsets])}
+        torch.segment_reduce(data, reduce, **segments).backward(grad_output)
+
+        self.assertEqual(data.grad, _amax_amin_reference_grad(reduce, grad_output))
+
+
+if __name__ == "__main__":
+    run_tests()

Automated by fix-implement.

Verify

Built pytorch origin/main@ec1459e58 from source with this branch as third_party/torch-xpu-ops, TORCH_XPU_ARCH_LIST=bmg,pvc, on Intel(R) Arc(TM) Pro B60 Graphics (BMG). The baseline and fixed builds differ only by the one hunk in SegmentReduceKernels.cpp.

Test case	Before	After
`reduce=max grad_out=-1 lengths`	FAIL	PASS
`reduce=max grad_out=-1 offsets`	FAIL	PASS
`reduce=min grad_out=-1 lengths`	FAIL	PASS
`reduce=min grad_out=-1 offsets`	FAIL	PASS

Before was segment_reduce=[-1.0, -1.0, 0.0] (sum -2.0) expected=[-0.5, -0.5, -0.0]; the issue's own reproducer now returns [-0.5, -0.5, 0.0], matching amax/amin backward exactly.

The usual caveat that a passing source build only proves "passes with the fix present" was closed out here: reverting only the SegmentReduceKernels.cpp hunk and rebuilding the same tree returned 4 failed / 4 passed (every negative-gradient combination failed, the positive-gradient half unaffected), and restoring the hunk returned all 8 to passing. The fix is what makes the test pass.

Regression checks on the fixed build:

Suite	Result
`test/regressions/test_segment_reduce_backward_ties.py`	8 passed
`test/xpu/test_segment_reductions_xpu.py`	74 passed
`test/xpu/test_ops_xpu.py -k segment_reduce`	34 passed, 2 skipped
`test/regressions/`	451 passed, 1 skipped, 1 failed

The single test/regressions/ failure is test_xpu_ops_header, which asserts on installed headers under torch/include/ATen/ops. An editable source build produces no torch/include at all (the generated headers stay in build/aten/src/ATen/ops/), so it fails for any such build and is unrelated to this change.

Verdict: PASSED -- the fix resolves all four reported cases and the attribution was confirmed by a revert-and-rebuild A/B
Lint: clean (lintrunner -a reported no issues and changed nothing; this repo ships no spin config, so lintrunner is its canonical linter)

Automated by fix-verify.

Summary

Outcome: IMPLEMENTING(fix_verified)

The MAX/MIN branch of the XPU segment_reduce backward kernel split a tied segment's gradient in two passes, but the second pass picked what to divide by the sign of the value it had just written rather than by tie membership, so a negative gradient was handed in full to each of the k tied elements and the segment's gradients summed to k * grad_output. The division pass now selects tied elements with the same isnan(value) || value == output predicate the assignment pass uses. A source build on BMG takes all four reported cases from FAIL to PASS and reproduces amax/amin backward exactly; reverting only that hunk on the same tree puts the four failures back, so the fix is demonstrably what resolves them, and the existing segment_reduce coverage (74 + 34 cases) is unchanged. A new XPU-only regression test covers tied segments with a negative grad_output across both the lengths= and offsets= entry points -- XPU-only because pytorch/pytorch#201633 reports the identical gate in the upstream CPU and CUDA kernels and is still open, so a device-generic version would fail on CPU.

Branch: agent/fix-issue-5808
Base: 73a6552c3c122614d2651de2a0c6feb8f81ae613
Changed files: src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp, test/regressions/test_segment_reduce_backward_ties.py
Nothing was pushed. The patch is exported as a workflow artifact for a human to review and apply.

One thing for the reviewer: SegmentReduceKernels.cpp:312-317 carries the if (idx >= size) { return; } guard twice, verbatim. It is pre-existing and untouched here.

Generated by fix job.

Review

Copy into a new comment, keep ONE verdict, add a note. Two verdicts in
one comment are dropped as ambiguous.

**strong-accepted** <!-- review: strong-accepted -->
**weak-accepted** <!-- review: weak-accepted -->
**weak-rejected** <!-- review: weak-rejected -->
**strong-rejected** <!-- review: strong-rejected -->
**issue-drop** <!-- review: issue-drop -->

notes:

On a batch run add item=<pytorch issue> to say which sub-item the
verdict is about, in the bold text and in the marker -- the marker
is invisible, so the bold line is all a reader sees. One comment per
sub-item, since a bare verdict on a batch does not say which test it
means and is dropped. run=<id> only when judging a run other than the
newest, in both places too.

A CI DISABLED sub-item, item= from the fan-out table:

**strong-accepted item=197334** <!-- review: strong-accepted item=197334 -->
notes: same TensorAssign overload as upstream CUDA, fp8 dispatch list matches

Another sub-item of the same batch, judged on an earlier run:

**weak-rejected run=35171979678 item=197335** <!-- review: weak-rejected run=35171979678 item=197335 -->
notes: green only because the test is skipped, the kernel still lacks fp8

An [xpu-alignment] issue that should not have been filed -- one issue
per difference, so no item=:

**issue-drop** <!-- review: issue-drop -->
notes: accepted XPU/CUDA numerics difference, nothing to repair

label-issue: #5808

Root cause: in segment_reduce_backward_kernel, the MAX/MIN tie-split loop divides by the tie count only when grad_input_data[data_index] > 0
(src/ATen/native/xpu/sycl/SegmentReduceKernels.cpp:359). A negative grad is never split; the fix belongs in torch-xpu-ops.

axis	value	reason
`type`	`Bug`	`issue_type` empty; reports an incorrect gradient (wrong result) from the XPU kernel.
`test`	`test: ut`	No e2e/oob signal; op-level reproducer, default UT surface.
`module`	`module: reduction`	Segment max/min reduction backward kernel `SegmentReduceKernels.cpp:290`.
`priority`	`Medium`	Incorrect result without a crash (functional error); single reproducer.
symptom	`Accuracy`	Silent wrong gradient: grad_sum -2.0 vs expected -1.0 (amax/amin reference).
symptom	`training`	The defect is in backward (gradient of segment_reduce).
triage	`duplicate`	Duplicate of pytorch/pytorch#201633 (MEDIUM, cross_link): same tie-split gate on the CPU/CUDA copies; the XPU kernel is a separate copy.

Applied to this issue. Generated by a local label-issue run.

Navigation Menu

Repository navigation

[xpu-alignment] [26-10-21] XPU segment_reduce max/min backward gives a negative tied gradient to every tied element instead of splitting it #5808

Summary

Upstream source

Observed XPU behavior

Target path and root cause

Reproducer

Environment

Ownership

Regression test

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[xpu-alignment] [26-10-21] XPU segment_reduce max/min backward gives a negative tied gradient to every tied element instead of splitting it #5808

Description

Summary

Upstream source

Observed XPU behavior

Target path and root cause

Reproducer

Environment

Ownership

Regression test

Activity

laifenxiawucha commented on Oct 28, 2026

torchxpubot commented on Oct 28, 2026

daisyden commented on Nov 10, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions