From ca120780cdf2ce4ca3f58cada288c97c3b3415c4 Mon Sep 17 00:00:00 2001
From: Aarav <211792885+AaravMehta-07@users.noreply.github.com>
Date: Mon, 12 Oct 2026 20:13:04 +0530
Subject: [PATCH] fix: use last axis when removing BOS tokens

---
 tests/unit/utilities/test_tokenize_utils.py  | 80 +++++++++++++++++++-
 transformer_lens/utilities/tokenize_utils.py |  2 +-
 2 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/tests/unit/utilities/test_tokenize_utils.py b/tests/unit/utilities/test_tokenize_utils.py
index 56d592d6eb..63596e07fc 100644
--- a/tests/unit/utilities/test_tokenize_utils.py
+++ b/tests/unit/utilities/test_tokenize_utils.py
@@ -1,12 +1,28 @@
-"""Tests for per-call padding-side overrides in tokenization utilities."""
+"""Tests for padding-side and sequence-axis tokenization semantics."""
 
 from copy import deepcopy
+from typing import cast
+from unittest.mock import Mock
 
+import pytest
 import torch
+from transformers import PreTrainedTokenizerBase
 
 from transformer_lens import utilities as utils
 
 
+def _tokenizer(
+    bos_token_id: int | None,
+    pad_token_id: int | None,
+    padding_side: str,
+) -> PreTrainedTokenizerBase:
+    tokenizer = Mock(spec=PreTrainedTokenizerBase)
+    tokenizer.bos_token_id = bos_token_id
+    tokenizer.pad_token_id = pad_token_id
+    tokenizer.padding_side = padding_side
+    return cast(PreTrainedTokenizerBase, tokenizer)
+
+
 def test_attention_mask_uses_explicit_padding_side(gpt2_tokenizer) -> None:
     tokenizer = deepcopy(gpt2_tokenizer)
     tokenizer.pad_token = tokenizer.eos_token
@@ -31,3 +47,65 @@ def test_bos_removal_uses_explicit_padding_side(gpt2_tokenizer) -> None:
     result = utils.get_tokens_with_bos_removed(tokenizer, tokens, padding_side="left")
 
     assert torch.equal(result, torch.tensor([[pad, 10, 11], [20, 21, 22]]))
+
+
+@pytest.mark.parametrize(
+    "tokens, expected, bos_token_id, pad_token_id, padding_side",
+    [
+        pytest.param([0, 1, 11, 12], [0, 11, 12], 1, 0, "left", id="rank-1-left"),
+        pytest.param(
+            [[0, 1, 11, 12], [1, 21, 22, 23]],
+            [[0, 11, 12], [21, 22, 23]],
+            1,
+            0,
+            "left",
+            id="rank-2-left-control",
+        ),
+        pytest.param([1], [], 1, 0, "left", id="bos-only-empty-result"),
+        pytest.param([0, 0, 11, 12], [0, 11, 12], 0, 0, "left", id="rank-1-bos-is-pad"),
+        pytest.param(
+            [[[0, 1, 11, 12], [0, 1, 21, 22]], [[1, 31, 32, 33], [0, 1, 41, 42]]],
+            [[[0, 11, 12], [0, 21, 22]], [[31, 32, 33], [0, 41, 42]]],
+            1,
+            0,
+            "left",
+            id="rank-3-left",
+        ),
+        pytest.param(
+            [[[0, 0, 0, 11, 12]], [[0, 21, 22, 23, 24]]],
+            [[[0, 0, 11, 12]], [[21, 22, 23, 24]]],
+            0,
+            0,
+            "left",
+            id="rank-3-bos-is-pad-leading-pads",
+        ),
+        pytest.param(
+            [[[1, 11, 12, 0], [1, 21, 22, 0]], [[1, 31, 32, 33], [1, 41, 42, 0]]],
+            [[[11, 12, 0], [21, 22, 0]], [[31, 32, 33], [41, 42, 0]]],
+            1,
+            0,
+            "right",
+            id="rank-3-right-control",
+        ),
+        pytest.param(
+            [[[0, 7, 8], [0, 9, 10]]],
+            [[[0, 7, 8], [0, 9, 10]]],
+            None,
+            0,
+            "left",
+            id="no-bos-unchanged",
+        ),
+    ],
+)
+def test_bos_removal_uses_last_axis(
+    tokens: list,
+    expected: list,
+    bos_token_id: int | None,
+    pad_token_id: int | None,
+    padding_side: str,
+) -> None:
+    result = utils.get_tokens_with_bos_removed(
+        _tokenizer(bos_token_id, pad_token_id, padding_side), torch.tensor(tokens)
+    )
+
+    assert torch.equal(result, torch.tensor(expected, dtype=torch.long))
diff --git a/transformer_lens/utilities/tokenize_utils.py b/transformer_lens/utilities/tokenize_utils.py
index 66b9459957..bcca9df56e 100644
--- a/transformer_lens/utilities/tokenize_utils.py
+++ b/transformer_lens/utilities/tokenize_utils.py
@@ -254,7 +254,7 @@ def get_tokens_with_bos_removed(
         else:
             real_bos_positions = (tokens == tokenizer.bos_token_id).int().argmax(-1)
 
-        tokens = tokens.scatter(dim=1, index=real_bos_positions.unsqueeze(-1), value=-100)
+        tokens = tokens.scatter(dim=-1, index=real_bos_positions.unsqueeze(-1), value=-100)
         return tokens[tokens != -100].view(*bos_removed_shape)