fix: simplify qkv_format checks using a membership tuple

[andrewwhitecdw]

This PR addresses the following issue in tests/pytorch/attention/run_attention_with_cp.py: simplify qkv_format checks using a membership tuple.

Changes

• tests/pytorch/attention/run_attention_with_cp.py: simplify qkv_format checks using a membership tuple.

Details

--- a/tests/pytorch/attention/run_attention_with_cp.py
+++ b/tests/pytorch/attention/run_attention_with_cp.py
@@ -1,2 +1,2 @@
-    if qkv_format == "bshd" or qkv_format == "sbhd":
-        seq_dim = qkv_format.index("s")
+    if qkv_format in ("bshd", "sbhd"):
+        seq_dim = qkv_format.index("s")

Tests

Let me know if you want tests added for this fix or not.

[greptile-apps]

Greptile Summary

This PR refactors two equivalent qkv_format conditionals in the context-parallel attention test helper to use tuple membership.

• Replaces repeated equality checks for bshd and sbhd without changing accepted formats or control flow.

Confidence Score: 5/5

The PR appears safe to merge because the refactor preserves the existing condition semantics and downstream behavior.

Both changed expressions accept exactly the same two string values as before and leave the associated sequence-dimension and comparison logic unchanged.

Important Files Changed

Filename	Overview
tests/pytorch/attention/run_attention_with_cp.py	Both changed conditions are behaviorally equivalent membership checks with no correctness, security, or test-coverage regression identified.

_{Reviews (1): Last reviewed commit: "fix: simplify qkv_format checks using a ..." | Re-trigger Greptile}