fix: remove redundant self-assignment out_ = out_

[andrewwhitecdw]

This PR addresses the following issue in tests/pytorch/attention/run_attention_with_cp.py: remove redundant self-assignment out_ = out_.

Changes

• tests/pytorch/attention/run_attention_with_cp.py: remove redundant self-assignment out_ = out_.

Details

--- a/tests/pytorch/attention/run_attention_with_cp.py
+++ b/tests/pytorch/attention/run_attention_with_cp.py
@@ -1,3 +1,2 @@
-        else:
-            out = out.index_select(0, seq_idx_q).contiguous()
-            out_ = out_
+        else:
+            out = out.index_select(0, seq_idx_q).contiguous()

Tests

• tests/pytorch/attention/test_redundant_self_assignment.py

--- /dev/null
+++ b/tests/pytorch/attention/test_redundant_self_assignment.py
@@ -0,0 +1,16 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Regression test for removal of redundant self-assignments."""
+
+import os
+
+
+def test_no_redundant_out_self_assignment():
+    """Ensure the forward-only THD branch no longer contains `out_ = out_`."""
+    file_path = os.path.join(os.path.dirname(__file__), "run_attention_with_cp.py")
+    with open(file_path, "r", encoding="utf-8") as f:
+        source = f.read()
+    assert "out_ = out_" not in source, (
+        "redundant self-assignment 'out_ = out_' must be removed"
+    )

[greptile-apps]

Greptile Summary

The PR removes a redundant context-parallel output self-assignment and avoids clearing the softmax-offset gradient during inference.

• Removes the no-op out_ = out_ statement from the forward-only THD comparison path.
• Restricts softmax-offset gradient clearing to training executions.
• Adds a CUDA-gated regression test for inference behavior.

Confidence Score: 3/5

The PR is not yet safe to merge because the new regression test fails when executed and its missing license header causes the pull-request license check to fail.

The test still constructs DotProductAttention with an unsupported softmax type, leaving softmax_offset undefined before the assertion, and the headerless file remains covered by the repository-wide license checker.

Files Needing Attention: tests/pytorch/attention/test_softmax_offset_inference.py

Important Files Changed

Filename	Overview
tests/pytorch/attention/run_attention_with_cp.py	Removes a redundant self-assignment and guards gradient clearing with is_training.
tests/pytorch/attention/test_softmax_offset_inference.py	Adds an inference regression test, but its unsupported softmax type prevents the intended assertion and its missing source header fails the license check.

_{Reviews (3): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

Unsupported softmax type breaks test

When this test runs on a CUDA-enabled system, softmax_type="softmax_offset" does not match any supported initialization branch, so softmax_offset is not initialized and the subsequent requires_grad access fails instead of testing eval behavior.

Suggested change

	core_attn = DotProductAttention(
	8, (64, 64), num_gqa_groups=4, softmax_type="softmax_offset"
	).cuda().eval()
	core_attn = DotProductAttention(
	8, (64, 64), num_gqa_groups=4, softmax_type="learnable"
	).cuda().eval()

Knowledge Base Used: Tests and QA

[greptile-apps]

Missing source license header

The new Python file omits the required NVIDIA copyright and license notices, causing the repository's L0 license checker to reject it.

Suggested change

	import pytest
	import torch
	from transformer_engine.pytorch import DotProductAttention
	# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	#
	# See LICENSE for license information.
	import pytest
	import torch
	from transformer_engine.pytorch import DotProductAttention

Knowledge Base Used: Tests and QA

[greptile-apps]

Regression test omitted from CI

The QA jobs enumerate test files explicitly, but none includes this new test, so CI silently skips the intended inference regression coverage.

Knowledge Base Used: Tests and QA

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[cyanguwa]

The changes in run_attention_with_cp.py looks good, but the test test_softmax_offset_inference.py‎ is probably unnecessary. The regular CI should test out the changes if any of them breaks anything. Thanks.