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 simplifies two equivalent qkv_format checks without changing their accepted values or branch behavior.

• Replaces repeated equality expressions with tuple membership checks.
• Applies the same refactor in context-parallel tensor preparation and result comparison.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
tests/pytorch/attention/run_attention_with_cp.py	Both modified conditions remain behaviorally equivalent, accepting only bshd and sbhd.

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