"""Mechanism tests use hand-written programs only, never experiment/final data.""" from dataclasses import replace from pathlib import Path from tempfile import TemporaryDirectory import pytest import torch from lightnest import rv32 from lightnest.relation_patch import RelationConfig, RelationPatchNetwork from lightnest.transformer_patch import ( VARIANTS, TransformerConfig, TransformerPatchNetwork, TransformerPredictor, attention_statistics, encode_batch, load_transformer, state_fingerprint, supervised_loss, ) @pytest.fixture(autouse=True) def small_cpu_budget(): torch.set_num_threads(2) def binary(words): return b''.join(word.to_bytes(4, 'little') for word in words) def fixtures(): first = binary([rv32.lw('x1', 0, 'x0'), rv32.lw('x2', 4, 'x0'), rv32.add('x3', 'x1', 'x2'), rv32.addi('x3', 'x3', 0), rv32.addi('x10', 'x3', 0), rv32.EBREAK]) short = binary([rv32.addi('x10', 'x0', 8), rv32.EBREAK]) return ['change add to sub', 'do not change sub to add'], [first, short] def inherited(variant='standard', seed=43, *, width=32): torch.manual_seed(7711) parent = RelationPatchNetwork(RelationConfig(width=width)) return TransformerPatchNetwork.from_relation( parent, config=TransformerConfig(variant=variant, width=width, ff_width=2 * width), seed=seed) def test_exact_common_inheritance_parameter_counts_roles_named_initialization_and_rng(): torch.manual_seed(7711) parent = RelationPatchNetwork(RelationConfig()).eval() rng = torch.get_rng_state().clone() children = [TransformerPatchNetwork.from_relation(parent, config=TransformerConfig(variant=v), seed=43) for v in VARIANTS] assert torch.equal(torch.get_rng_state(), rng) counts = dict(standard=146930, relation_bias=146946, reverse_query=146930, shared_recurrent=113458, differential=147122, centered_attention=146938) reference = dict(children[0][0].named_parameters()) for child, report in children: assert child.parameter_count == report['parameters'] == counts[child.config.variant] assert report['copied_parameter_values'] == 79858 assert report['newly_initialized_parameter_values'] == child.parameter_count - 79858 assert report['all_copied_tensors_exact'] and not report['function_preserving'] parameters = dict(child.named_parameters()) assert set(report['parameter_name_map']) == set(parameters) == set(child.parameter_roles()) assert sum(child.parameter_counts_by_role().values()) == child.parameter_count for name, source in report['parameter_name_map'].items(): if source is not None: old = dict(parent.named_parameters())[source] assert torch.equal(parameters[name], old) assert parameters[name].data_ptr() != old.data_ptr() if name in reference: assert torch.equal(parameters[name], reference[name]), name roles = child.parameter_roles() assert roles['command_conv1.weight'] == roles['command_conv2.weight'] == 'auxiliary' assert roles['command_embedding.weight'] == 'embedding' assert roles['pointer.weight'] == 'output' assert roles['blocks.0.attention.q_proj.weight'] == 'hidden_matrix' assert all(parameters[name].ndim == 2 for name, role in roles.items() if role == 'hidden_matrix') repeated, _ = TransformerPatchNetwork.from_relation(parent, config=child.config, seed=43) assert state_fingerprint(repeated) == state_fingerprint(child) different, _ = TransformerPatchNetwork.from_relation(parent, config=child.config, seed=44) assert not torch.equal(different.blocks[0].attention.q_proj.weight, child.blocks[0].attention.q_proj.weight) assert torch.equal(different.command_fc.weight, child.command_fc.weight) def test_standard_attention_is_scaled_qk_softmax_and_separate_qkv(): model, _ = inherited() encoded = encode_batch(*fixtures(), model.config) _, _, diag = model(*encoded, return_diagnostics=True) block = model.blocks[0] h = block.norm1(diag['initial_states']) batch, length, width = h.shape heads, dim = model.config.heads, width // model.config.heads q = block.attention.q_proj(h).reshape(batch, length, heads, dim).transpose(1, 2) k = block.attention.k_proj(h).reshape(batch, length, heads, dim).transpose(1, 2) logits = q @ k.transpose(-1, -2) / dim**.5 expected = logits.masked_fill(~diag['valid'][:, None, None, :], -torch.inf).softmax(-1) expected = expected.masked_fill(~diag['valid'][:, None, :, None], 0.) actual = diag['attention'][0] torch.testing.assert_close(actual['qk_logits'], logits, rtol=0, atol=0) torch.testing.assert_close(actual['probabilities'], expected, rtol=0, atol=0) pointers = [getattr(block.attention, name).weight.data_ptr() for name in ('q_proj', 'k_proj', 'v_proj')] assert len(set(pointers)) == 3 assert not any(isinstance(module, torch.nn.GRUCell) for module in model.modules()) @pytest.mark.parametrize('variant', VARIANTS) def test_all_variants_have_finite_supervised_gradients_and_masked_probabilities(variant): model, _ = inherited(variant) commands, codes = fixtures() scores, bits, diag = model(*encode_batch(commands, codes, model.config), return_diagnostics=True) assert scores.shape == (2, 1025) and bits.shape == (2, 1024, 8) assert bool(torch.isfinite(scores).all() and torch.isfinite(bits).all()) assert bool((scores[1, len(codes[1]):model.none_index] == -1e9).all()) assert int(torch.count_nonzero(bits[1, len(codes[1]):])) == 0 for stage in diag['attention']: for field in ('probabilities', 'secondary_probabilities'): if field not in stage: continue weights = stage[field] expected = stage['query_valid'][:, None].expand(-1, weights.shape[1], -1).to(weights.dtype) torch.testing.assert_close(weights.sum(-1), expected, atol=1e-6, rtol=1e-6) assert int(torch.count_nonzero(weights.masked_select(~stage['key_valid'][:, None, None, :]))) == 0 rows = [{'edits': [{'offset': 11, 'value': codes[0][11] ^ 64}]}, {'edits': []}] loss = supervised_loss(scores, bits, rows, model.none_index)[0] loss.backward() for block in model.blocks: for layer in (block.attention.q_proj, block.attention.k_proj, block.attention.v_proj, block.attention.out_proj, block.ff_in, block.ff_out): grad = layer.weight.grad assert grad is not None and bool(torch.isfinite(grad).all()) assert float(grad.abs().max()) > 1e-10 assert all(value.grad is None or bool(torch.isfinite(value.grad).all()) for value in model.parameters()) extra = dict(relation_bias='relation_bias', differential='lambda_q1', centered_attention='suppression_logits').get(variant) if extra is not None: assert float(getattr(model.blocks[0].attention, extra).grad.abs().max()) > 1e-10 stats = attention_statistics(diag) assert len(stats['stages']) == model.config.steps @pytest.mark.parametrize('variant', VARIANTS) def test_padding_batch_invariance_and_no_edges_outside_explicit_bias_arm(variant): model, _ = inherited(variant) model.eval() commands, codes = fixtures() encoded = encode_batch(commands, codes, model.config) before = state_fingerprint(model) with torch.inference_mode(): together = model(*encoded) solo = model(*encode_batch(commands[1:], codes[1:], model.config)) n = encoded[2].shape[-1] extended_command = torch.nn.functional.pad(encoded[0], (0, 3)) extended_binary = torch.nn.functional.pad(encoded[1], (0, 8)) extended_edges = torch.nn.functional.pad(encoded[2], (0, 2, 0, 2)) extended = model(extended_command, extended_binary, extended_edges) if variant != 'relation_bias': altered = model(encoded[0], encoded[1], torch.rand_like(encoded[2])) for original, changed in zip(together, altered): torch.testing.assert_close(original, changed, atol=0, rtol=0) for combined, single, padded in zip(together, solo, extended): torch.testing.assert_close(combined[1], single[0], atol=2e-6, rtol=2e-5) torch.testing.assert_close(combined, padded, atol=2e-6, rtol=2e-5) assert state_fingerprint(model) == before def test_direct_relation_bias_zero_is_standard_and_nonzero_adds_only_direct_edges(): standard, _ = inherited('standard') biased, _ = inherited('relation_bias') encoded = encode_batch(*fixtures(), standard.config) with torch.inference_mode(): for original, zero_bias in zip(standard(*encoded), biased(*encoded)): torch.testing.assert_close(original, zero_bias, atol=0, rtol=0) biased.blocks[0].attention.relation_bias[0].fill_(2.5) _, _, diag = biased(*encoded, return_diagnostics=True) stage = diag['attention'][0] edges = stage['typed_direct_edges'] assert int(torch.count_nonzero(edges.diagonal(dim1=-2, dim2=-1))) == 0 assert float(edges[0, 0, 4, 3]) == 1. # x10 copy reads the immediately previous x3 definition. assert float(edges[0, 0, 4, 2]) == 0. # No transitive closure through its x3 copy. expected = 2.5 * edges[:, 0, None].expand(-1, biased.config.heads, -1, -1) torch.testing.assert_close(stage['biased_logits'] - stage['qk_logits'], expected, atol=1e-6, rtol=1e-6) def test_shared_block_really_reuses_weights_and_reverse_query_is_last_valid_token(): commands, codes = fixtures() shared, _ = inherited('shared_recurrent') calls = [] hook = shared.blocks[0].register_forward_hook(lambda module, args, result: calls.append(id(module))) with torch.inference_mode(): default = shared(*encode_batch(commands, codes, shared.config), return_diagnostics=True) single = shared(*encode_batch(commands, codes, shared.config), steps=1) hook.remove() assert len(shared.blocks) == 1 and len(calls) == 5 and len(set(calls)) == 1 assert default[2]['compute']['executed_blocks'] == 4 assert not torch.equal(default[0], single[0]) reverse, _ = inherited('reverse_query') with torch.inference_mode(): _, _, diag = reverse(*encode_batch(commands, codes, reverse.config), return_diagnostics=True) assert diag['return_query_index'].tolist() == [5, 1] torch.testing.assert_close(diag['return_query_initial'], diag['stage_states'][0][torch.arange(2), [5, 1]]) assert diag['attention'][0]['probabilities'].shape == (2, 4, 6, 6) assert diag['attention'][1]['probabilities'].shape == (2, 4, 1, 6) assert diag['compute']['dense_attention_score_elements'] == 2 * 4 * (6 * 6 + 6) assert default[2]['compute']['dense_attention_score_elements'] == 4 * 2 * 4 * 6 * 6 def test_differential_and_centered_subtractions_report_signed_weights_truthfully(): differential, _ = inherited('differential') with torch.no_grad(): differential.blocks[0].attention.lambda_q1.fill_(.5) differential.blocks[0].attention.lambda_k1.fill_(.5) differential.blocks[0].attention.lambda_q2.zero_() differential.blocks[0].attention.lambda_k2.zero_() with torch.inference_mode(): _, _, diag = differential(*encode_batch(*fixtures(), differential.config), return_diagnostics=True) stage = diag['attention'][0] assert stage['effective_heads'] == 2 and stage['softmax_streams'] == 2 assert stage['probabilities'].shape == (2, 2, 6, 6) torch.testing.assert_close(stage['effective_weights'], stage['probabilities'] - stage['coefficient'] * stage['secondary_probabilities'], atol=0, rtol=0) assert float(stage['coefficient']) > 1 and bool((stage['effective_weights'] < 0).any()) stats = attention_statistics(diag) assert min(stats['stages'][0]['effective_negative_mass_per_head']) > 0 assert max(stats['stages'][0]['effective_row_sum_per_head']) < 0 centered, _ = inherited('centered_attention') with torch.inference_mode(): _, _, diag = centered(*encode_batch(*fixtures(), centered.config), return_diagnostics=True) stage = diag['attention'][0] uniform = stage['key_valid'].float() / stage['key_valid'].sum(1, keepdim=True) expected = stage['probabilities'] - .5 * uniform[:, None, None] expected = expected.masked_fill(~stage['query_valid'][:, None, :, None], 0.) torch.testing.assert_close(stage['effective_weights'], expected, atol=0, rtol=0) @pytest.mark.parametrize('variant', VARIANTS) def test_checkpoint_roundtrip_keeps_tensor_inventory_config_roles_and_rng(variant): model, _ = inherited(variant) with TemporaryDirectory(prefix='lightnest-transformer-test-') as directory: for suffix in ('.pt', '.pt.gz'): path = Path(directory) / ('model' + suffix) model.save(path, purpose='mechanism test only') rng = torch.get_rng_state().clone() loaded = load_transformer(path) assert torch.equal(torch.get_rng_state(), rng) assert loaded.config == model.config assert loaded.initialization_report == model.initialization_report assert loaded.parameter_roles() == model.parameter_roles() assert state_fingerprint(loaded) == state_fingerprint(model) with torch.inference_mode(): encoded = encode_batch(*fixtures(), model.config) for actual, expected in zip(loaded(*encoded), model(*encoded)): torch.testing.assert_close(actual, expected, atol=0, rtol=0) def test_loader_rejects_nonfinite_or_wrong_shape_and_predictor_uses_expected_byte(): from lightnest.iterative_patch import run_iterative_patch model, _ = inherited() with TemporaryDirectory(prefix='lightnest-transformer-invalid-') as directory: path = Path(directory) / 'invalid.pt' model.save(path) saved = torch.load(path, weights_only=True) saved['state_dict']['blocks.0.attention.q_proj.weight'][0, 0] = float('nan') torch.save(saved, path) with pytest.raises(ValueError, match='tensor'): load_transformer(path) model.save(path) saved = torch.load(path, weights_only=True) saved['state_dict']['blocks.0.attention.q_proj.weight'] = torch.zeros(1) torch.save(saved, path) with pytest.raises(ValueError, match='tensor'): load_transformer(path) class FixedNetwork: config = TransformerConfig() none_index = 1024 def eval(self): return self def __call__(self, command, code, relations): scores = torch.full((1, 1025), -40.) bits = torch.full((1, 1024, 8), -40.) position = 1024 if int(code[0, 3]) - 1 == 64 else 3 scores[0, position] = 40. bits[0, 3, 6] = 40. return scores, bits code = binary([rv32.add('x3', 'x1', 'x2'), rv32.addi('x10', 'x3', 0), rv32.EBREAK]) predictor = TransformerPredictor(FixedNetwork()) decision = predictor('change add to sub', code) assert decision.edits[0].expected == code[3:4] result = run_iterative_patch('change add to sub', code, predictor, max_steps=2) assert result.final[3] == 64 and result.final[:3] == code[:3] and result.final[4:] == code[4:] assert result.stop_reason == 'predictor_stop' def test_declared_capacity_covers_512_byte_program_and_invalid_config_rejected(): model, _ = inherited() code = binary([rv32.addi('x3', 'x3', 0)] * 126 + [rv32.addi('x10', 'x3', 0), rv32.EBREAK]) assert len(code) == 512 with torch.inference_mode(): scores, bits = model(*encode_batch(['change add to sub'], [code], model.config)) assert scores.shape == (1, 1025) and bits.shape == (1, 1024, 8) assert bool(torch.isfinite(scores).all() and torch.isfinite(bits).all()) with pytest.raises(ValueError, match='even'): TransformerConfig(variant='differential', heads=1) with pytest.raises(ValueError, match='divide'): TransformerConfig(heads=3) with pytest.raises(ValueError, match='query'): TransformerConfig(variant='reverse_query', steps=3) with pytest.raises(ValueError, match='runtime'): model(*encode_batch(*fixtures(), model.config), steps=4)