"""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()