From df52a2268f2d77a38e8f4481e47c0ef2559187a2 Mon Sep 17 00:00:00 2001 From: Alessandro Potenza Date: Sun, 23 Aug 2026 16:44:44 +0200 Subject: [PATCH] fix: pin the eager decoder in the AuxK test so it runs anywhere The test builds CPU tensors and its docstring says it uses the eager decoder fallback, but nothing made that true: `decoder_impl` is bound at import time and prefers the Triton kernel, which rejects CPU tensors. The test therefore fails on any machine where Triton imports, GPU or not, so `pytest` is red on a fresh clone. Patch `decoder_impl` for the duration of the test and call `eager_decode` for the reference value. --- tests/test_auxk_loss.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/test_auxk_loss.py b/tests/test_auxk_loss.py index 71b58bcd..6b9852a2 100644 --- a/tests/test_auxk_loss.py +++ b/tests/test_auxk_loss.py @@ -1,16 +1,19 @@ import torch +import sparsify.sparse_coder as sparse_coder_module from sparsify import SparseCoder, SparseCoderConfig -from sparsify.utils import decoder_impl +from sparsify.utils import eager_decode -def test_auxk_loss_does_not_double_count_b_dec(): +def test_auxk_loss_does_not_double_count_b_dec(monkeypatch): """The AuxK loss target ``e = y - sae_out`` already accounts for ``b_dec`` (since ``sae_out`` includes it), so the second decoder pass used to compute ``e_hat`` must *not* add ``b_dec`` again. See issue #132. - This runs on CPU using the eager decoder fallback, so it requires no GPU. + This runs on CPU, so it needs the eager decoder: `decoder_impl` is bound at + import time and prefers the Triton kernel, which rejects CPU tensors. """ + monkeypatch.setattr(sparse_coder_module, "decoder_impl", eager_decode) torch.manual_seed(0) d_in = 16 @@ -49,7 +52,7 @@ def test_auxk_loss_does_not_double_count_b_dec(): # Correct target: decode without adding b_dec a second time. assert sae.W_dec is not None - e_hat = decoder_impl(auxk_indices, auxk_acts.to(sae.dtype), sae.W_dec.mT) + e_hat = eager_decode(auxk_indices, auxk_acts.to(sae.dtype), sae.W_dec.mT) expected_auxk_loss = scale * (e_hat - e.detach()).pow(2).sum() / total_variance torch.testing.assert_close(out.auxk_loss, expected_auxk_loss)