diff --git a/src/steerbench/vectors.py b/src/steerbench/vectors.py index a57e5c0..b29c1ed 100644 --- a/src/steerbench/vectors.py +++ b/src/steerbench/vectors.py @@ -123,6 +123,10 @@ def normalize_alpha( Returns all three units. See :func:`dose` for the scalar full dose. """ + if layer not in vector.directions: + raise ValueError( + f"layer {layer} not in steering vector; available layers: {sorted(vector.directions)}" + ) norm = float(torch.linalg.vector_norm(vector.directions[layer])) by_vector = alpha * norm by_residual = by_vector / residual_norm if residual_norm is not None else None diff --git a/tests/test_vectors.py b/tests/test_vectors.py index d094bb7..816f240 100644 --- a/tests/test_vectors.py +++ b/tests/test_vectors.py @@ -212,3 +212,21 @@ def test_gguf_superset_loads_in_repeng(tmp_path: Path) -> None: assert hint == "gpt2" names = {t.name for t in reader.tensors} assert names == {"direction.5", "direction.6"} + + +def test_normalize_alpha_out_of_range_layer_names_available() -> None: + vec = SteeringVector(directions={3: torch.tensor([1.0, 0.0]), 7: torch.tensor([0.0, 2.0])}) + with pytest.raises(ValueError, match=r"layer 999 not in steering vector.*\[3, 7\]"): + normalize_alpha(vec, layer=999, alpha=1.0) + + +def test_dose_out_of_range_layer_names_available() -> None: + vec = SteeringVector(directions={3: torch.tensor([1.0, 0.0]), 7: torch.tensor([0.0, 2.0])}) + with pytest.raises(ValueError, match=r"layer 999 not in steering vector.*\[3, 7\]"): + dose(vec, layer=999, alpha=1.0, residual_norm=10.0) + + +def test_normalize_alpha_valid_layer_unchanged_by_guard() -> None: + vec = SteeringVector(directions={3: torch.tensor([3.0, 4.0])}) # norm 5 + got = normalize_alpha(vec, layer=3, alpha=2.0) + assert got.by_vector_norm == pytest.approx(10.0)