Spaces:
Running on Zero
Running on Zero
File size: 1,559 Bytes
9d24374 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | from __future__ import annotations
import torch
from featurelens.sae import SAEWeights
def make_sae() -> SAEWeights:
w_enc_t = torch.tensor(
[
[1.0, 0.0, -1.0, 0.5, 0.0],
[0.0, 1.0, 0.0, 0.5, -1.0],
[0.0, 0.0, 1.0, 0.0, 1.0],
]
)
w_dec = torch.tensor(
[
[1.0, 0.0, 0.0, 0.5, 0.0],
[0.0, 1.0, 0.0, 0.5, 0.0],
[0.0, 0.0, 1.0, 0.0, 1.0],
]
)
return SAEWeights(
layer=0,
w_enc_t=w_enc_t,
w_dec=w_dec,
b_enc=torch.zeros(5),
b_dec=torch.zeros(3),
top_k=2,
)
def test_encode_keeps_topk_relu_features() -> None:
sae = make_sae()
encoding = sae.encode(torch.tensor([2.0, 1.0, -1.0]))
assert encoding.indices.shape == (2,)
assert set(encoding.indices.tolist()) == {0, 3}
assert encoding.active_count == 2
assert encoding.activation_for(0) == 2.0
assert encoding.activation_for(4) == 0.0
def test_decode_sparse_uses_selected_decoder_columns() -> None:
sae = make_sae()
encoding = sae.encode(torch.tensor([2.0, 1.0, -1.0]))
reconstructed = sae.decode_sparse(encoding)
assert reconstructed.shape == (3,)
assert torch.isfinite(reconstructed).all()
def test_batched_encode_and_decode() -> None:
sae = make_sae()
hidden = torch.tensor([[2.0, 1.0, -1.0], [0.0, 2.0, 2.0]])
encoding = sae.encode(hidden)
decoded = sae.decode_sparse(encoding)
assert encoding.indices.shape == (2, 2)
assert decoded.shape == hidden.shape
|