From cb10cab89407f97218a446975bd1af5a433944e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:52:46 +0200 Subject: [PATCH] [JAX] Fix typo in qkv_layout docstring: seperate -> separate --- transformer_engine/jax/attention.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index ecca4a3871..c7cb733d4c 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -120,11 +120,11 @@ class QKVLayout(Enum): BSHD Format: - BS3HD: q,k,v are interleave packed as a tensor with shape [b, s, 3, h, d]. - BSHD_BS2HD: q with shape [b, s, h, d] and kv are interleaved with shape [b, s, 2, h, d]. - - BSHD_BSHD_BSHD: q,k,v are seperate tensors with shape [b, s, h, d] + - BSHD_BSHD_BSHD: q,k,v are separate tensors with shape [b, s, h, d] THD Format: Shape is same as BSHD layout but allow multiple segments packed in a sequence. - T3HD: q,k,v are interleave packed as a tensor with shape [b, s, 3, h, d]. - THD_T2HD: q with shape [b, s, h, d] and kv are interleaved with shape [b, s, 2, h, d]. - - THD_THD_THD: q,k,v are seperate tensors with shape [b, s, h, d] + - THD_THD_THD: q,k,v are separate tensors with shape [b, s, h, d] """ BS3HD = NVTE_QKV_Layout.NVTE_BS3HD