File size: 7,085 Bytes
e64c5d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """Checks for the new architectural connection, action contract, and checkpoints."""
from dataclasses import replace
import json
import pytest
import torch
from lightnest import rv32
from lightnest.relation_patch import (
RelationConfig, RelationPatchNetwork, RelationPredictor, definition_use_graph,
encode_batch, load_records, load_relation, supervised_loss,
)
def assemble(words):
return b''.join(x.to_bytes(4, 'little') for x in words)
def far_pair():
words = [rv32.lw('x1', 0, 'x0'), rv32.lw('x2', 4, 'x0'), rv32.add('x3', 'x1', 'x2')]
words += [rv32.addi('x0', 'x0', 31 + i) for i in range(8)]
q = len(words) * 4 + 3
words += [rv32.add('x4', 'x1', 'x2')]
words += [rv32.addi('x0', 'x0', 211 + i) for i in range(8)]
pair = [assemble(words + [rv32.addi('x10', reg, 0), rv32.EBREAK]) for reg in ('x3', 'x4')]
return pair, (11, q)
def test_definition_edges_obey_overwrites_zero_and_implicit_return():
code = assemble([rv32.lw('x1', 0, 'x0'), rv32.lw('x2', 4, 'x0'),
rv32.add('x3', 'x1', 'x2'), rv32.sub('x3', 'x2', 'x1'),
rv32.addi('x0', 'x3', 0), rv32.addi('x10', 'x3', 0), rv32.EBREAK])
graph = definition_use_graph(code)
assert graph.instruction_count == 7
assert set(graph.edges) == {(2, 0), (2, 1), (3, 0), (3, 1), (4, 3), (5, 3), (6, 5)}
assert (5, 2) not in graph.edges
@pytest.mark.parametrize('code', [b'', b'\x13', assemble([rv32.beq('x1', 'x2', 4)]),
assemble([rv32.EBREAK, rv32.addi('x0', 'x0', 0)])])
def test_unsupported_structure_is_explicit(code):
with pytest.raises(ValueError):
definition_use_graph(code)
def test_local_pointer_is_invariant_but_graph_pointer_reads_remote_return():
torch.set_num_threads(2)
torch.manual_seed(137)
pair, (p, q) = far_pair()
# Exactly identical 17-byte neighborhoods at both candidate positions.
for pos in (p, q):
assert pair[0][pos - 8:pos + 9] == pair[1][pos - 8:pos + 9]
graph = RelationPatchNetwork(RelationConfig(mode='dataflow')).eval()
local = RelationPatchNetwork(replace(graph.config, mode='local')).eval()
local.load_state_dict(graph.state_dict())
assert sum(p.numel() for p in graph.parameters()) == sum(p.numel() for p in local.parameters())
encoded = encode_batch(['change the returned add to sub'] * 2, pair, graph.config)
with torch.inference_mode():
ls, _ = local(*encoded)
gs, _ = graph(*encoded)
local_change = (ls[0, p] - ls[0, q]) - (ls[1, p] - ls[1, q])
graph_change = (gs[0, p] - gs[0, q]) - (gs[1, p] - gs[1, q])
assert float(local_change.abs()) < 1e-7
# This proves a connection exists, not that random weights solve the task.
assert float(graph_change.abs()) > 1e-6
def test_dynamic_batch_padding_preserves_inference_and_deeper_passes():
torch.manual_seed(17)
model = RelationPatchNetwork(RelationConfig(width=32)).eval()
pair, _ = far_pair()
short = assemble([rv32.addi('x10', 'x0', 9), rv32.EBREAK])
with torch.inference_mode():
first = model(*encode_batch(['no changes'], [short], model.config))
together = model(*encode_batch(['no changes', 'replace the returned add with sub'], [short, pair[0]], model.config))
deeper = model(*encode_batch(['no changes'], [short], model.config), message_steps=8)
for x, y in zip(first, together):
torch.testing.assert_close(x[0], y[0], atol=2e-6, rtol=2e-5)
assert deeper[0].shape == (1, 1025)
assert not torch.equal(first[0], deeper[0])
def test_set_target_loss_accepts_both_orders_and_explicit_stop():
scores = torch.zeros(2, 9, requires_grad=True)
bits = torch.zeros(2, 8, 8, requires_grad=True)
rows = [{'edits': [{'offset': 1, 'value': 65}, {'offset': 5, 'value': 2}]}, {'edits': []}]
loss, pointer_loss, byte_loss = supervised_loss(scores, bits, rows, none_index=8)
expected = (-torch.log(torch.tensor(2 / 9)) - torch.log(torch.tensor(1 / 9))) / 2
torch.testing.assert_close(pointer_loss, expected)
torch.testing.assert_close(byte_loss, torch.tensor(2.).log())
loss.backward()
assert scores.grad[0, 1] < 0 and scores.grad[0, 5] < 0 and scores.grad[1, 8] < 0
assert torch.count_nonzero(bits.grad[0, 1]) == 8
assert torch.count_nonzero(bits.grad[0, 5]) == 8
assert torch.count_nonzero(bits.grad[1]) == 0
def test_checkpoint_roundtrip_and_reject_wrong_architecture_or_nonfinite(tmp_path):
model = RelationPatchNetwork(RelationConfig(width=32, mode='local'))
path = tmp_path / 'model.pt'
model.save(path)
restored = load_relation(path)
assert restored.config == model.config
for a, b in zip(model.parameters(), restored.parameters()):
torch.testing.assert_close(a, b, rtol=0., atol=0.)
saved = torch.load(path, weights_only=True)
saved['architecture'] = 'absolute-byte-control'
torch.save(saved, path)
with pytest.raises(ValueError, match='architecture'):
load_relation(path)
saved['architecture'] = model.architecture
saved['state_dict']['pointer.weight'][0, 0] = float('nan')
torch.save(saved, path)
with pytest.raises(ValueError, match='tensor'):
load_relation(path)
def test_adapter_uses_actual_byte_precondition_and_explicit_stop():
from lightnest.iterative_patch import run_iterative_patch
class FixedNetwork:
config = RelationConfig()
none_index = 1024
def eval(self):
return self
def __call__(self, command, binary, relations, message_steps=None):
scores = torch.full((1, 1025), -40.)
bits = torch.full((1, 1024, 8), -40.)
# A fixture action tests the adapter only; this is not a trained model.
value = int(binary[0, 3]) - 1
if value == 64:
scores[0, 1024] = 40.
else:
scores[0, 3] = 40.
bits[0, 3, 6] = 40.
return scores, bits
code = assemble([rv32.add('x3', 'x1', 'x2'), rv32.addi('x10', 'x3', 0), rv32.EBREAK])
predictor = RelationPredictor(FixedNetwork())
decision = predictor('change add to sub', code)
assert decision.edits[0].offset == 3
assert decision.edits[0].expected == code[3:4]
assert decision.edits[0].replacement == bytes([64])
result = run_iterative_patch('change add to sub', code, predictor, max_steps=2)
assert result.final[3] == 64
assert result.final[:3] == code[:3] and result.final[4:] == code[4:]
def test_dataset_checks_do_not_silently_miss_an_edit(tmp_path):
binary = assemble([rv32.addi('x10', 'x0', 9), rv32.EBREAK])
path = tmp_path / 'rows.jsonl'
row = {'id': 'case', 'command': 'do not change', 'binary': binary.hex(), 'target': binary.hex(), 'edits': []}
path.write_text(json.dumps(row) + '\n')
assert load_records(path)[0]['edits'] == []
row['target'] = bytes([binary[0] ^ 1]).hex() + binary[1:].hex()
path.write_text(json.dumps(row) + '\n')
with pytest.raises(ValueError, match='disagree'):
load_records(path)
|