happy lint

Signed-off-by: Michael Goldfarb <[email protected]>

transformer_engine/jax/cpp_extensions/attention.py

@@ -946,7 +946,7 @@ def ag(x):
             case NVTE_QKV_Layout.NVTE_BSHD_BSHD_BSHD:
                 return ag(k), ag(v)
             case _:
-                raise ValueError(f"Unsupported layout {self.qkv_layout=}")
+               return k, v
 
     def reduce_scatter_dkv(self, dk, dv):
         """Performs a reduce-scatter of dk and dv over context parallel ranks."""
@@ -960,7 +960,7 @@ def rs(x):
             case NVTE_QKV_Layout.NVTE_BSHD_BSHD_BSHD:
                 return rs(dk), rs(dv)
             case _:
-                raise ValueError(f"Unsupported layout {self.qkv_layout=}")
+                return dk, dv
 
     def kv_seqlens_for_rank(self, cp_rank, kv_max_seqlen, kv_seqlen_per_subrank):
         """Returns sequence lengths of KV to use for each sub rank of the given cp_rank.
@@ -1001,7 +1001,7 @@ def sliced(x):
             case NVTE_QKV_Layout.NVTE_BSHD_BSHD_BSHD:
                 return sliced(k), sliced(v)
             case _:
-                raise ValueError(f"Unsupported layout {self.qkv_layout=}")
+                return k, v
 
     def pad_kv(self, dk, dv, pad_seq_len):
         """Pads dk and dv tensors to a sequence length of pad_seq_len."""
@@ -1017,7 +1017,7 @@ def pad(x, npad):
                 npad = [[0, 0], [0, pad_seq_len], [0, 0], [0, 0]]
                 return pad(dk, npad), pad(dv, npad)
             case _:
-                raise ValueError(f"Unsupported layout {self.qkv_layout=}")
+                return dk, dv
 
 
 class FusedAttnCPWithAllGatherFwdPrimitive(FusedAttnFwdPrimitive):