tests/test_gemma_scope.py

"""Gemma Scope (JumpReLU) SAE support: npz loading, gating, and interventions.

Download-free: a synthetic npz in the Gemma Scope naming convention
(layer_*/width_*/average_l0_*/params.npz with W_enc, W_dec, b_enc, b_dec,
threshold arrays) is planted in a temp directory. The fetch helper is never
run against the network here.
"""

import tempfile
import unittest
from pathlib import Path

import numpy as np
import torch

from owmi.interventions import (
    apply_sae_feature_intervention,
    fetch_gemma_scope_sae,
    gemma_scope_param_path,
    load_linear_sae,
    make_forward_hook,
)
from owmi.types import InterventionSpec
from owmi.validation import validate_sae_roundtrip

D_MODEL = 4
D_SAE = 3
# Per-feature JumpReLU thresholds with known planted structure.
THRESHOLD = np.array([0.5, 1.5, 0.5], dtype=np.float32)


def _write_gemma_scope_npz(root, layer=20, width="16k", average_l0=91, **overrides):
    """Synthesize a params.npz in the Gemma Scope layout and array orientation.

    W_enc is stored [d_model, d_sae] and W_dec [d_sae, d_model], matching the
    real google/gemma-scope-9b-it-res files. Encoder/decoder are the first
    D_SAE coordinates of the identity, so feature j reads and writes channel j
    and the pre-activation of feature j equals x[..., j] - b_dec[j] + b_enc[j].
    """
    arrays = {
        "W_enc": np.eye(D_MODEL, dtype=np.float32)[:, :D_SAE],
        "W_dec": np.eye(D_MODEL, dtype=np.float32)[:D_SAE, :],
        "b_enc": np.zeros(D_SAE, dtype=np.float32),
        "b_dec": np.zeros(D_MODEL, dtype=np.float32),
        "threshold": THRESHOLD,
    }
    arrays.update(overrides)
    rel = gemma_scope_param_path(layer, width, average_l0)
    path = Path(root) / rel
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez(path, **arrays)
    return str(path)


class GemmaScopeLoadingTests(unittest.TestCase):
    def setUp(self):
        tmpdir = tempfile.TemporaryDirectory()
        self.addCleanup(tmpdir.cleanup)
        self.root = tmpdir.name
        self.npz_path = _write_gemma_scope_npz(self.root)

    def test_npz_autodetected_as_jumprelu_with_normalized_shapes(self):
        sae = load_linear_sae(self.npz_path)
        self.assertEqual(sae.activation, "jumprelu")
        self.assertEqual((sae.d_model, sae.d_sae), (D_MODEL, D_SAE))
        self.assertEqual(tuple(sae.W_enc.shape), (D_MODEL, D_SAE))
        self.assertEqual(tuple(sae.W_dec.shape), (D_SAE, D_MODEL))
        self.assertTrue(torch.equal(sae.threshold, torch.from_numpy(THRESHOLD)))

    def test_npz_missing_array_raises(self):
        bad_dir = Path(self.root) / "bad"
        bad_dir.mkdir()
        bad_path = str(bad_dir / "params.npz")
        np.savez(
            bad_path,
            W_enc=np.eye(D_MODEL, dtype=np.float32)[:, :D_SAE],
            b_enc=np.zeros(D_SAE, dtype=np.float32),
            b_dec=np.zeros(D_MODEL, dtype=np.float32),
        )
        with self.assertRaises(ValueError):
            load_linear_sae(bad_path)

    def test_npz_without_threshold_defaults_to_relu(self):
        path = _write_gemma_scope_npz(self.root, average_l0=14)
        # Rewrite without the threshold array.
        with np.load(path) as arrays:
            keep = {k: arrays[k] for k in ("W_enc", "W_dec", "b_enc", "b_dec")}
        np.savez(path, **keep)
        sae = load_linear_sae(path)
        self.assertEqual(sae.activation, "relu")
        self.assertIsNone(sae.threshold)

    def test_explicit_relu_flag_overrides_npz_threshold(self):
        sae = load_linear_sae(self.npz_path, activation="relu")
        self.assertEqual(sae.activation, "relu")
        self.assertIsNone(sae.threshold)
        x = torch.tensor([[1.0, 1.0, 0.0, 7.0]])
        # ReLU keeps the sub-threshold feature 1 that JumpReLU would gate out.
        self.assertTrue(torch.equal(sae.encode(x), torch.tensor([[1.0, 1.0, 0.0]])))

    def test_torch_state_dict_still_loads_as_relu(self):
        pt_path = str(Path(self.root) / "sae.pt")
        torch.save({
            "W_enc": torch.eye(D_MODEL),
            "b_enc": torch.zeros(D_MODEL),
            "W_dec": torch.eye(D_MODEL),
            "b_dec": torch.zeros(D_MODEL),
        }, pt_path)
        sae = load_linear_sae(pt_path)
        self.assertEqual(sae.activation, "relu")
        self.assertIsNone(sae.threshold)

    def test_torch_state_dict_with_threshold_loads_as_jumprelu(self):
        pt_path = str(Path(self.root) / "sae_jump.pt")
        torch.save({
            "W_enc": torch.eye(D_MODEL)[:, :D_SAE],
            "b_enc": torch.zeros(D_SAE),
            "W_dec": torch.eye(D_MODEL)[:D_SAE, :],
            "b_dec": torch.zeros(D_MODEL),
            "threshold": torch.from_numpy(THRESHOLD),
        }, pt_path)
        sae = load_linear_sae(pt_path)
        self.assertEqual(sae.activation, "jumprelu")
        self.assertTrue(torch.equal(sae.threshold, torch.from_numpy(THRESHOLD)))

    def test_threshold_length_mismatch_raises(self):
        path = _write_gemma_scope_npz(
            self.root, average_l0=25, threshold=np.zeros(D_SAE + 1, dtype=np.float32)
        )
        with self.assertRaises(ValueError):
            load_linear_sae(path)


class JumpReLUGatingTests(unittest.TestCase):
    def setUp(self):
        tmpdir = tempfile.TemporaryDirectory()
        self.addCleanup(tmpdir.cleanup)
        self.sae = load_linear_sae(_write_gemma_scope_npz(tmpdir.name))

    def test_below_threshold_encodes_to_zero_above_passes_pre_activation(self):
        # pre_act = [1.0, 1.0, 0.6]; thresholds [0.5, 1.5, 0.5].
        x = torch.tensor([[1.0, 1.0, 0.6, 7.0]])
        features = self.sae.encode(x)
        self.assertEqual(float(features[0, 0]), 1.0)   # above 0.5: passes through
        self.assertEqual(float(features[0, 1]), 0.0)   # below 1.5: gated to zero
        self.assertAlmostEqual(float(features[0, 2]), 0.6, places=6)  # above 0.5: passes through

    def test_gate_is_strict_inequality_at_threshold(self):
        x = torch.tensor([[0.5, 1.5, 0.5, 0.0]])  # pre_act exactly at threshold
        self.assertTrue(torch.equal(self.sae.encode(x), torch.zeros((1, D_SAE))))

    def test_negative_pre_activation_is_gated_unlike_relu_zero(self):
        x = torch.tensor([[-2.0, 3.0, 0.0, 0.0]])
        features = self.sae.encode(x)
        self.assertEqual(float(features[0, 0]), 0.0)
        self.assertEqual(float(features[0, 1]), 3.0)

    def test_decode_roundtrip_recovers_above_threshold_channels(self):
        x = torch.tensor([[1.0, 2.0, 0.6, 7.0]])  # all features above threshold
        recon = self.sae.decode(self.sae.encode(x))
        # Channels 0..2 are reconstructed; channel 3 is residual (b_dec = 0).
        self.assertTrue(torch.equal(recon[0, :D_SAE], x[0, :D_SAE]))
        self.assertEqual(float(recon[0, 3]), 0.0)


class JumpReLUInterventionTests(unittest.TestCase):
    """Delta-form substitution semantics on a Gemma Scope-format SAE."""

    def setUp(self):
        tmpdir = tempfile.TemporaryDirectory()
        self.addCleanup(tmpdir.cleanup)
        self.sae_path = _write_gemma_scope_npz(tmpdir.name)
        self.sae = load_linear_sae(self.sae_path)
        # Positions 0/2 have feature 0 above threshold (1.0, 3.0); position 1 below (0.2).
        self.hidden = torch.tensor([
            [[1.0, 2.0, 0.6, 7.0],
             [0.2, 2.0, 0.6, 7.0],
             [3.0, 2.0, 0.6, 7.0]],
        ])

    def _spec(self, mode="zero", **kwargs):
        defaults = dict(
            layer_index=0, mode=mode, object_kind="sae_feature",
            feature_id=0, sae_weights_path=self.sae_path,
        )
        defaults.update(kwargs)
        return InterventionSpec(**defaults)

    def test_zero_changes_only_feature_channel_and_skips_gated_positions(self):
        out = apply_sae_feature_intervention(self.hidden, self._spec(), self.sae)
        # Above-threshold positions: coefficient (== channel 0 here) removed.
        self.assertEqual(float(out[0, 0, 0]), 0.0)
        self.assertEqual(float(out[0, 2, 0]), 0.0)
        # Below-threshold position: f_j = 0 already, so zeroing is a no-op.
        self.assertTrue(torch.equal(out[0, 1], self.hidden[0, 1]))
        # Off-feature channels (including the residual channel 3) untouched everywhere.
        self.assertTrue(torch.equal(out[..., 1:], self.hidden[..., 1:]))

    def test_delta_form_moves_only_along_decoder_direction(self):
        out = apply_sae_feature_intervention(self.hidden, self._spec(mode="scale", strength=3.0), self.sae)
        delta = out - self.hidden
        direction = self.sae.W_dec[0]
        coeff = self.sae.encode(self.hidden)[..., 0]
        expected = (coeff * 3.0 - coeff).unsqueeze(-1) * direction
        self.assertTrue(torch.allclose(delta, expected))

    def test_scale_at_strength_one_is_exact_identity_for_jumprelu(self):
        out = apply_sae_feature_intervention(self.hidden, self._spec(mode="scale", strength=1.0), self.sae)
        self.assertTrue(torch.equal(out, self.hidden))

    def test_token_positions_restrict_the_intervention(self):
        spec = self._spec(token_positions=[2])
        out = apply_sae_feature_intervention(self.hidden, spec, self.sae, position_seq_len=3)
        self.assertTrue(torch.equal(out[0, :2], self.hidden[0, :2]))
        self.assertEqual(float(out[0, 2, 0]), 0.0)

    def test_replace_baseline_transfers_reference_coefficient(self):
        reference = self.hidden.clone()
        reference[..., 0] = 5.0  # above threshold everywhere
        spec = self._spec(mode="replace", reference_mode="baseline")
        out = apply_sae_feature_intervention(self.hidden, spec, self.sae, reference_hidden=reference)
        # Delta form: x' = x + (f'_j - f_j) W_dec[j]. At above-threshold
        # positions f_j equals channel 0, so x'[0] becomes exactly 5.0; at the
        # gated position f_j = 0 and the sub-threshold residual 0.2 is
        # preserved on top of the transferred coefficient (0.2 + 5.0).
        expected = self.hidden[..., 0] + (5.0 - self.sae.encode(self.hidden)[..., 0])
        self.assertTrue(torch.allclose(out[..., 0], expected))
        self.assertEqual(float(out[0, 0, 0]), 5.0)
        self.assertAlmostEqual(float(out[0, 1, 0]), 5.2, places=5)
        self.assertTrue(torch.equal(out[..., 1:], self.hidden[..., 1:]))

    def test_forward_hook_routes_jumprelu_sae_prefill_only(self):
        hook = make_forward_hook(self._spec(), prompt_length=3, sae=self.sae)
        modified = hook(torch.nn.Identity(), (), (self.hidden,))[0]
        self.assertEqual(float(modified[0, 0, 0]), 0.0)
        decode_step = torch.ones((1, 1, D_MODEL))
        unchanged = hook(torch.nn.Identity(), (), (decode_step,))[0]
        self.assertTrue(torch.equal(unchanged, decode_step))

    def test_roundtrip_validation_accepts_npz_path_and_reports_jumprelu(self):
        report = validate_sae_roundtrip(self.sae_path, self.hidden[0], self._spec())
        self.assertEqual(report["activation"], "jumprelu")
        self.assertTrue(report["feature_delta_confirmed"])
        self.assertTrue(report["delta_in_feature_direction"])


class FetchHelperTests(unittest.TestCase):
    """Download-free checks only; fetch_gemma_scope_sae never hits the network here."""

    def test_param_path_matches_hub_convention(self):
        self.assertEqual(
            gemma_scope_param_path(20, "16k", 91),
            "layer_20/width_16k/average_l0_91/params.npz",
        )
        self.assertEqual(
            gemma_scope_param_path(9, "131k", 121),
            "layer_9/width_131k/average_l0_121/params.npz",
        )

    def test_fetch_returns_existing_file_without_downloading(self):
        tmpdir = tempfile.TemporaryDirectory()
        self.addCleanup(tmpdir.cleanup)
        rel = gemma_scope_param_path(31, "16k", 76)
        existing = _write_gemma_scope_npz(tmpdir.name, layer=31, width="16k", average_l0=76)
        result = fetch_gemma_scope_sae(
            "google/gemma-scope-9b-it-res", rel, tmpdir.name
        )
        self.assertEqual(result, existing)
        # And the returned path is directly loadable.
        self.assertEqual(load_linear_sae(result).activation, "jumprelu")


if __name__ == "__main__":
    unittest.main()