"""Numerical and presentation contracts for analysis results on the CPU.""" from types import SimpleNamespace as NS import matplotlib.pyplot as plt import numpy as np import pytest import torch from obliteratus.analysis import visualization as plots from obliteratus.analysis.spectral_certification import SpectralCertifier from obliteratus.analysis.tuned_lens import RefusalTunedLens, TunedLensProbe pytestmark = pytest.mark.cpu @pytest.fixture(autouse=True) def close_figures(): yield plt.close("all") @pytest.mark.parametrize("save", [False, True]) def test_token_spectrum_preserves_order_signs_and_layer_selection(tmp_path, monkeypatch, save): shown = [] monkeypatch.setattr(plt, "show", lambda: shown.append(True)) layer = NS( layer_idx=7, top_promoted=[("yes", 3.0), ("perhaps", 1.0)], top_suppressed=[("no", -4.0), ("never", -2.0)], refusal_compliance_gap=2.5, refusal_specificity=0.75, ) result = NS(per_layer={7: layer}, strongest_refusal_layer=7) path = tmp_path / "tokens.png" if save else None fig = plots.plot_logit_lens_spectrum( result, layer_idx=7 if save else None, output_path=path, title="Spectrum" if save else None ) ax = fig.axes[0] assert [patch.get_width() for patch in ax.patches] == [-2, -4, 3, 1] assert [label.get_text() for label in ax.get_yticklabels()] == [ "'never'", "'no'", "'yes'", "'perhaps'", ] assert "2.5000" in ax.texts[0].get_text() assert ax.get_title() == ("Spectrum" if save else "Logit Lens — Layer 7") if save: assert path.read_bytes().startswith(b"\x89PNG") assert not shown else: assert shown == [True] assert plots.plot_logit_lens_spectrum(result, layer_idx=99) is None @pytest.mark.parametrize("save", [False, True]) def test_pareto_averages_scores_and_preserves_reference_coordinates(tmp_path, monkeypatch, save): monkeypatch.setattr(plt, "show", lambda: None) path = tmp_path / "pareto.png" if save else None fig = plots.plot_capability_safety_pareto( {"a": NS(score=0.2), "b": NS(score=0.8)}, 0.3, other_points=[(0.9, 0.6, "Baseline")], output_path=path, ) assert np.asarray(fig.axes[0].collections[0].get_offsets()).tolist() == [[0.3, 0.5]] assert np.asarray(fig.axes[0].collections[1].get_offsets()).tolist() == [[0.9, 0.6]] assert "Baseline" in [text.get_text() for text in fig.axes[0].texts] if save: assert path.read_bytes().startswith(b"\x89PNG") empty = plots.plot_capability_safety_pareto({}, 0.0) assert np.asarray(empty.axes[0].collections[0].get_offsets()).tolist() == [[0.0, 0.0]] def test_topology_missing_means_have_zero_strength_and_squeezes_direction(monkeypatch): monkeypatch.setattr(plt, "show", lambda: None) fig = plots.plot_refusal_topology( {4: torch.tensor([[3.0, 0.0]]), 8: torch.tensor([0.0, 1.0])}, {4: torch.tensor([2.0, 7.0])}, {4: torch.tensor([0.0, 7.0])}, [4], ) assert [p.get_height() for p in fig.axes[0].patches] == [2.0, 0.0] assert [t.get_text() for t in fig.axes[0].get_xticklabels()] == ["4", "8"] @pytest.mark.parametrize( "path", ["embed_out", "output", "transformer.wte", "model.embed_tokens", "gpt_neox.embed_in"] ) def test_tuned_lens_finds_supported_output_or_tied_embeddings(path): weight = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) node = NS(weight=weight) for part in reversed(path.split(".")): node = NS(**{part: node}) torch.testing.assert_close(RefusalTunedLens()._get_unembedding_matrix(node), weight) def test_tuned_lens_missing_output_is_actionable(): with pytest.raises(RuntimeError, match="Cannot locate unembedding"): RefusalTunedLens()._get_unembedding_matrix(NS()) def test_tuned_lens_group_filters_invalid_and_unencodable_tokens(): def encode(text, **kwargs): if text == "error": raise ValueError("unknown token") return {"valid": [1, 0], "empty": [], "negative": [-1], "outside": [3]}[text] boosts = RefusalTunedLens()._get_token_group_boosts( torch.tensor([2.0, 5.0, 9.0]), NS(encode=encode), ["error", "empty", "negative", "outside", "valid"], ) assert boosts == [5.0] def test_tuned_lens_affine_direction_and_report_have_known_oracle(): lens = RefusalTunedLens(top_k=1) tokenizer = NS( decode=lambda ids: str(ids[0]), encode=lambda text, **kw: [0] if text == "sorry" else [] ) model = NS( lm_head=NS(weight=torch.eye(2)), model=NS(norm=NS(weight=torch.tensor([2.0, 3.0]), bias=torch.tensor([1.0, -1.0]))), ) probe = TunedLensProbe(5, torch.eye(2), torch.tensor([100.0, 100.0]), 0.0) result = lens.analyze_all_layers( {5: torch.tensor([[1.0, 0.0]]), 6: torch.ones(2)}, {5: probe}, model, tokenizer ) layer = result.per_layer[5] assert layer.top_promoted == [("0", 3.0)] assert layer.top_suppressed == [("1", -1.0)] assert layer.refusal_compliance_gap == 3.0 assert layer.correction_magnitude == 0.0 report = lens.format_report(result) assert "Strongest refusal layer: 5" in report assert "Mean refusal-compliance gap: 3.0000" in report assert "Top promoted:" in report assert "Top suppressed:" in report empty = lens.analyze_all_layers({6: torch.ones(2)}, {}, model, tokenizer) assert empty.per_layer == {} assert empty.mean_refusal_compliance_gap == 0.0 assert "No layers analyzed." in lens.format_report(empty) def test_tuned_lens_comparison_intersects_layers_and_detects_reversed_ranking(): result = NS( per_layer={ 1: NS(refusal_compliance_gap=1.0), 2: NS(refusal_compliance_gap=2.0), 3: NS(refusal_compliance_gap=3.0), } ) assert RefusalTunedLens.compare_with_logit_lens(result, {1: 1.0}) == 1.0 assert ( RefusalTunedLens.compare_with_logit_lens(result, {1: 3.0, 2: 2.0, 3: 1.0, 4: 100.0}) == -1.0 ) def test_spectral_diagonal_oracle_and_degenerate_condition_estimates(): certifier = SpectralCertifier() covariance = torch.diag(torch.tensor([0.0, 2.0, 8.0])) result = certifier._eigenvalue_analysis(covariance, bbp_threshold=4.0, mp_upper=2.0) torch.testing.assert_close(result.eigenvalues, torch.tensor([8.0, 2.0, 0.0])) torch.testing.assert_close( covariance @ result.eigenvectors, result.eigenvectors @ torch.diag(result.eigenvalues) ) assert result.above_threshold == [0] assert result.in_bulk == [1] assert result.signal_subspace_dim == 1 assert certifier._estimate_condition_number(covariance) == 4.0 assert certifier._estimate_condition_number(torch.zeros(2, 2)) == 1.0 assert certifier._estimate_noise_variance(covariance, n=2, d=3) == 2.0 ratio = (1 - (3 / 12) ** 0.5) ** 2 + (3 / 12) ** (1 / 3) assert certifier._estimate_noise_variance(covariance, n=12, d=3) == pytest.approx(2 / ratio) assert certifier._estimate_noise_variance(torch.zeros(2, 2), n=2, d=2) == 1e-10 for spectrum in [ torch.tensor([]), torch.tensor([1.0]), torch.tensor([1.0, float("nan")]), torch.tensor([-1.0, 2.0]), ]: assert certifier._estimate_condition_number_from_spectrum(spectrum) == 1.0 assert certifier._estimate_condition_number_from_spectrum(torch.tensor([1.0, 1e8])) == 1e6 def test_spectral_linear_algebra_failure_returns_conservative_fallbacks(monkeypatch): def fail(*args, **kwargs): raise RuntimeError("eigensolver failed") monkeypatch.setattr(torch.linalg, "eigvalsh", fail) monkeypatch.setattr(torch.linalg, "eigh", fail) certifier = SpectralCertifier() assert certifier._estimate_noise_variance(torch.eye(2), 2, 2) == 1.0 assert certifier._estimate_condition_number(torch.eye(2)) == 1.0 result = certifier._eigenvalue_analysis(torch.eye(2), 1.0, 1.0) assert result.signal_subspace_dim == 0 assert result.above_threshold == [] assert result.in_bulk == [] torch.testing.assert_close(result.eigenvalues, torch.zeros(1))