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 guards softmax-offset gradient clearing so it occurs only during training and removes a redundant output self-assignment. It also adds a focused eval-mode regression test.

Confidence Score: 3/5

The PR is not yet safe to merge because the regression test fails on CUDA systems and the new file fails the mandatory license check.

The test still supplies an unsupported softmax type, leaving softmax_offset undefined before its assertion, while the headerless test file is rejected by the pull-request license workflow.

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

Important Files Changed

Filename	Overview
tests/pytorch/attention/run_attention_with_cp.py	Restricts gradient clearing to training runs and removes the no-op out_ self-assignment.
tests/pytorch/attention/test_softmax_offset_inference.py	Adds a CUDA-gated eval-mode regression test for softmax-offset gradient state.

_{Reviews (2): 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.