"""Narrow mutation/inheritance tests; no experimental datasets are read.""" from pathlib import Path from tempfile import TemporaryDirectory import pytest import torch from lightnest import rv32 from lightnest.morph_patch import ( MorphConfig, MorphPatchNetwork, MorphPredictor, compare_functions, encode_batch, load_morph, route_statistics, state_fingerprint, supervised_loss, ) from lightnest.relation_patch import RelationConfig, RelationPatchNetwork @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 test_one_expert_exactly_inherits_relation_without_rng_or_storage_alias(): torch.manual_seed(1801) parent = RelationPatchNetwork(RelationConfig(width=64)).eval() state = torch.get_rng_state().clone() child, report = MorphPatchNetwork.from_relation(parent, seed=81) assert torch.equal(state, torch.get_rng_state()) assert child.parameter_count == 113138 assert report['all_copied_tensors_exact'] and report['newly_initialized_parameter_values'] == 0 assert report['parameter_name_map']['experts.0.weight_ih'] == 'update.weight_ih' assert child.experts[0].weight_ih.data_ptr() != parent.update.weight_ih.data_ptr() result = compare_functions(parent, child, *fixtures(), atol=0, rtol=0) assert result['passed'] and result['scores_max_abs_diff'] == result['bits_max_abs_diff'] == 0 def test_selected_single_leaf_splits_preserve_function_and_old_tensor_names(): torch.manual_seed(13) parent = MorphPatchNetwork().eval() commands, codes = fixtures() for index, expected in zip((0, 0, 1), (138227, 163316, 188405)): before = state_fingerprint(parent) rng = torch.get_rng_state().clone() child, report = parent.spawn_child('split', expert_index=index, seed=271 + index) assert child.parameter_count == expected assert report['parameter_name_map'][f'experts.{parent.config.experts}.weight_hh'] == f'experts.{index}.weight_hh' assert report['newly_initialized_parameter_values'] == 129 assert report['function_preserving_expected'] assert torch.equal(torch.get_rng_state(), rng) assert state_fingerprint(parent) == before assert compare_functions(parent, child, commands, codes)['passed'] assert child.experts[-1].weight_ih.data_ptr() != parent.experts[index].weight_ih.data_ptr() parent = child assert parent.config.router_children == ((1, 2), (-1, -3), (-2, -4)) with pytest.raises(ValueError, match='fewer than four'): parent.spawn_child('split', expert_index=0) def test_new_state_gates_break_gradient_symmetry_without_initial_output_noise(): torch.manual_seed(31) one = MorphPatchNetwork(MorphConfig(width=32)) two, _ = one.spawn_child('split', expert_index=0, seed=79) commands, codes = fixtures() encoded = encode_batch(commands, codes, two.config) scores, bits, diagnostics = two(*encoded, return_diagnostics=True) initial = route_statistics(diagnostics, two) assert initial['route_input_dependence_nats'] > 0 assert initial['expert_parameter_rms_distance'][0][1] == 0 assert initial['expert_output_rms_distance'][0][1] == 0 rows = [{'edits': [{'offset': 11, 'value': codes[0][11] ^ 64}]}, {'edits': []}] optimizer = torch.optim.SGD(two.parameters(), lr=0.1) loss, _, _ = supervised_loss(scores, bits, rows, two.none_index) loss.backward() difference = (two.experts[0].weight_ih.grad - two.experts[1].weight_ih.grad).abs().max() assert float(difference) > 1e-8 optimizer.step() _, _, after = two(*encoded, return_diagnostics=True) measured = route_statistics(after, two) assert measured['expert_parameter_rms_distance'][0][1] > 0 assert measured['expert_output_rms_distance'][0][1] > 0 optimizer.zero_grad(set_to_none=True) scores, bits = two(*encoded) supervised_loss(scores, bits, rows, two.none_index)[0].backward() assert float(two.routers[0].weight.grad.abs().max()) > 1e-10 def test_padding_diagnostics_uniform_ablation_and_deeper_weight_inheritance(): torch.manual_seed(43) model = MorphPatchNetwork(MorphConfig(width=32, experts=3)).eval() commands, codes = fixtures() encoded = encode_batch(commands, codes, model.config) before = state_fingerprint(model) with torch.inference_mode(): scores, bits, diag = model(*encoded, return_diagnostics=True) uniform_scores, _, uniform = model(*encoded, return_diagnostics=True, routing='uniform') solo = model(*encode_batch(commands[1:], codes[1:], model.config)) mask = diag['valid'][None].expand(diag['route_weights'].shape[:3]) torch.testing.assert_close(diag['route_weights'].sum(-1)[mask], torch.ones(int(mask.sum()))) assert int(torch.count_nonzero(diag['route_weights'][~mask])) == 0 torch.testing.assert_close(uniform['route_weights'][mask], torch.full((int(mask.sum()), 3), 1 / 3)) assert float((scores - uniform_scores).abs().max()) > 1e-7 assert state_fingerprint(model) == before for together, single in zip((scores, bits), solo): torch.testing.assert_close(together[1], single[0], atol=2e-6, rtol=2e-5) deeper, report = model.spawn_child('deepen', message_steps=8, seed=2) assert not report['function_preserving_expected'] assert state_fingerprint(deeper) == before and deeper.parameter_count == model.parameter_count assert not compare_functions(model, deeper, commands, codes, atol=0, rtol=0)['passed'] assert compare_functions(model, deeper, commands, codes, message_steps=4, atol=0, rtol=0)['passed'] @pytest.mark.parametrize('tree', [((0, -1),), ((-1, -1),), ((-1, -3),), ((1, -2), (-1, -2))]) def test_invalid_router_topology_rejected(tree): with pytest.raises(ValueError): MorphConfig(experts=len(tree) + 1, router_children=tree) def test_checkpoint_roundtrip_preserves_lineage_and_rejects_bad_tensors(): torch.manual_seed(59) model = MorphPatchNetwork(MorphConfig(width=32)) model, _ = model.spawn_child('split', expert_index=0, seed=3) model, _ = model.spawn_child('split', expert_index=1, seed=5) with TemporaryDirectory(prefix='lightnest-morph-test-') as directory: path = Path(directory) / 'model.pt' model.save(path, purpose='mechanism test only') restored = load_morph(path) assert restored.config == model.config and restored.lineage == model.lineage assert state_fingerprint(restored) == state_fingerprint(model) assert compare_functions(model, restored, *fixtures(), atol=0, rtol=0)['passed'] saved = torch.load(path, weights_only=True) saved['state_dict']['routers.0.weight'][0, 0] = float('nan') torch.save(saved, path) with pytest.raises(ValueError, match='tensor'): load_morph(path) def test_predictor_byte_precondition_and_stop_use_common_runner(): from lightnest.iterative_patch import run_iterative_patch class FixedNetwork: config = MorphConfig() none_index = 1024 def eval(self): return self def __call__(self, command, code, relations, message_steps=None, routing='learned'): 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 = MorphPredictor(FixedNetwork(), routing='uniform') 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'