| import copy,random,unittest |
| import torch |
| from torch.nn import functional as F |
| from lightnest.gradient_interference_experiment import RoutedXorPolicy,freeze,teacher_metrics |
| from lightnest.xor_patch_experiment import inputs,XorPolicy |
| from lightnest.multifamily_patch import make_task,teacher_states |
|
|
| class GradientRoutingTests(unittest.TestCase): |
| def setUp(self): |
| torch.set_num_threads(2);torch.manual_seed(17) |
| self.states=teacher_states(make_task(random.Random(31))) |
| def test_byte_gradient_isolation(self): |
| for detached in [False,True]: |
| m=RoutedXorPolicy(detached);c,b,p,y=inputs(self.states);mask=p<256 |
| _,bits=m(c,b);F.binary_cross_entropy_with_logits(bits[torch.arange(len(p))[mask],p[mask]],y[mask]).backward() |
| backbone=sum(float(v.grad.abs().sum()) for k,v in m.named_parameters() if not k.startswith('replacement.') and v.grad is not None) |
| self.assertGreater(float(m.replacement.weight.grad.abs().sum()),0) |
| if detached:self.assertEqual(backbone,0) |
| else:self.assertGreater(backbone,0) |
| def test_pointer_gradient_unaffected(self): |
| a=RoutedXorPolicy();b=copy.deepcopy(a);b.detach_byte=True;c,x,p,y=inputs(self.states) |
| for m in [a,b]:F.cross_entropy(m(c,x)[0],p).backward() |
| self.assertGreater(float(a.fusion.weight.grad.abs().sum()),0) |
| for (k,v),(l,w) in zip(a.named_parameters(),b.named_parameters()): |
| if v.grad is None:self.assertIsNone(w.grad) |
| else:self.assertTrue(torch.equal(v.grad,w.grad),k) |
| def test_inference_equivalent_to_existing_xor(self): |
| a=RoutedXorPolicy();b=XorPolicy(copy.deepcopy(a));c,x,_,_=inputs(self.states) |
| |
| from lightnest.learned_patch import PatchNetwork |
| net=PatchNetwork();net.load_state_dict(a.state_dict());b=XorPolicy(net) |
| for lhs,rhs in zip(a(c,x),b(c,x)):self.assertTrue(torch.equal(lhs,rhs)) |
| a.detach_byte=True |
| for lhs,rhs in zip(a(c,x),b(c,x)):self.assertTrue(torch.equal(lhs,rhs)) |
| def test_frozen_disjoint_and_reproducible(self): |
| train,suites=freeze();train2,suites2=freeze();self.assertEqual(train,train2);self.assertEqual(suites,suites2) |
| self.assertFalse({(t.command,t.binary) for t in train}&{(t.command,t.binary) for ts in suites.values() for t in ts}) |
| if __name__=='__main__':unittest.main() |
|
|
