def test_place_ensembles(): with nengo.Network() as net: add_params(net) offchip = nengo.Ensemble(10, 1, label="offchip") net.config[offchip].on_chip = False direct = nengo.Ensemble(1, 1, neuron_type=nengo.Direct(), label="direct") with nengo.Network(): onchip = nengo.Ensemble(20, 1, label="onchip") pre = nengo.Ensemble(10, 1, label="pre") post = nengo.Ensemble(10, 1, label="post") error = nengo.Ensemble(10, 1, label="error") conn = nengo.Connection(pre, post, learning_rule_type=nengo.PES()) nengo.Connection(error, conn.learning_rule) networks = SplitNetworks(net, node_neurons=default_node_neurons) place_ensembles(networks) assert networks.moves[offchip] == "host" assert networks.moves[direct] == "host" assert networks.moves[onchip] == "chip" assert networks.moves[pre] == "chip" assert networks.moves[post] == "host" assert networks.moves[error] == "host"
def test_place_probes(): with nengo.Network() as net: offchip1 = nengo.Node(0) with nengo.Network(): onchip1 = nengo.Ensemble(10, 1) offchip2 = nengo.Ensemble(10, 1) onchip2 = nengo.Ensemble(10, 1) onchip3 = nengo.Connection(onchip1, onchip2) offchip3 = nengo.Connection(offchip1, offchip2) offchip_probes = [ nengo.Probe(offchip1), nengo.Probe(offchip2), nengo.Probe(offchip3), ] onchip_probes = [ nengo.Probe(onchip1), nengo.Probe(onchip2), nengo.Probe(onchip3), ] networks = SplitNetworks(net, node_neurons=default_node_neurons) for obj in [offchip1, offchip2, offchip3]: networks.move(obj, "host") for obj in [onchip1, onchip2, onchip3]: networks.move(obj, "chip") place_probes(networks) assert all(networks.location(p) == "host" for p in offchip_probes) assert all(networks.location(p) == "chip" for p in onchip_probes)
def test_splitnetwork_remove_add(): net = nengo.Network() networks = SplitNetworks(net) e = nengo.Ensemble(1, 1, add_to_container=False) networks.add(e, "chip") networks.remove(e) assert e not in networks.adds
def test_place_nodes(): with nengo.Network() as net: offchip1 = nengo.Node(0) with nengo.Network(): offchip2 = nengo.Node(np.sin) offchip3 = HostSendNode(dimensions=1) onchip = ChipReceiveNode(dimensions=1, size_out=1) networks = SplitNetworks(net, node_neurons=default_node_neurons) place_nodes(networks) assert networks.moves[offchip1] == "host" assert networks.moves[offchip2] == "host" assert networks.moves[offchip3] == "host" assert networks.moves[onchip] == "chip"
def test_split_host_to_chip_no_node_neurons(): with nengo.Network() as net: node_offchip = nengo.Node(np.sin) ens_onchip = nengo.Ensemble(10, 1) conn = nengo.Connection(node_offchip, ens_onchip) networks = SplitNetworks(net) networks.move(node_offchip, "host") networks.move(ens_onchip, "chip") with pytest.raises(BuildError): split_host_to_chip(networks, conn)
def test_place_inter_network_connection(): with nengo.Network() as net: offchip = nengo.Ensemble(10, 1) onchip = nengo.Ensemble(10, 1) onon = nengo.Connection(onchip, onchip) onoff = nengo.Connection(onchip, offchip) offon = nengo.Connection(offchip, onchip) offoff = nengo.Connection(offchip, offchip) networks = SplitNetworks(net, node_neurons=default_node_neurons) networks.move(onchip, "chip") networks.move(offchip, "host") place_internetwork_connections(networks, networks.original.all_connections) assert onoff not in networks assert offon not in networks assert networks.location(onon) == "chip" assert networks.location(offoff) == "host"
def test_split_host_neurons_to_chip(): with nengo.Network() as net: offchip = nengo.Ensemble(10, 1) onchip = nengo.Ensemble(10, 1) neurons2neurons = nengo.Connection(offchip.neurons, onchip.neurons, transform=np.ones((10, 10))) neurons2ensemble = nengo.Connection(offchip.neurons, onchip, transform=np.ones((1, 10))) networks = SplitNetworks(net, node_neurons=default_node_neurons) networks.move(offchip, "host") networks.move(onchip, "chip") def assert_split_correctly(split_conn): assert len(networks.adds) == 4 added_types = sorted([(type(obj).__name__, location) for obj, location in networks.adds.items()]) assert added_types == [ ("ChipReceiveNeurons", "chip"), ("Connection", "chip"), ("Connection", "host"), ("HostSendNode", "host"), ] assert split_conn in networks.removes send = next(obj for obj in networks.adds if isinstance(obj, HostSendNode)) receive = next(obj for obj in networks.adds if isinstance(obj, ChipReceiveNeurons)) assert networks.host2chip_senders[send] is receive split_host_neurons_to_chip(networks, neurons2neurons) assert_split_correctly(neurons2neurons) networks.adds.clear() # Makes testing subsequent adds easier split_host_neurons_to_chip(networks, neurons2ensemble) assert_split_correctly(neurons2ensemble)
def test_split_pre_from_host(): with nengo.Network() as net: pre_1 = nengo.Node(0, label="pre_1") pre_2 = nengo.Ensemble(10, 1, label="pre_2") pre_3 = nengo.Node(size_in=1, label="pre_3") pre_4 = nengo.Ensemble(1, 1, label="pre_4") send = HostSendNode(dimensions=1) onchip = nengo.Ensemble(1, 1, label="onchip") post1 = nengo.Ensemble(10, 1, label="post1") post2 = nengo.Node(size_in=1, label="post2") pre_connections = [ nengo.Connection(pre_1, pre_2), nengo.Connection(pre_2, pre_3), nengo.Connection(pre_3, pre_4), nengo.Connection(pre_4.neurons, send), ] post_connections = [ nengo.Connection(onchip, post1), nengo.Connection(post1, post2), ] networks = SplitNetworks(net, node_neurons=default_node_neurons) for obj in [pre_1, pre_3, send, post1, post2]: networks.move(obj, "host") for obj in [pre_2, pre_4]: networks.add(obj, "host") for conn in pre_connections + post_connections: networks.move(conn, "host") networks.move(onchip, "chip") split_pre_from_host(networks) for obj in [pre_1, pre_2, pre_3, pre_4, send] + pre_connections: assert networks.location(obj) == "host_pre", obj for obj in [post1, post2] + post_connections: assert networks.location(obj) == "host", obj assert networks.location(onchip) == "chip"
def test_split_host_to_learning_rule(): with nengo.Network() as net: pre = nengo.Ensemble(10, 1, label="pre") post = nengo.Ensemble(10, 1, label="post") err_onchip = nengo.Ensemble(10, 1, label="err_onchip") err_offchip = nengo.Ensemble(10, 1, label="err_offchip") ens_conn = nengo.Connection(pre, post, learning_rule_type=nengo.PES()) neurons_conn = nengo.Connection(pre.neurons, post.neurons, learning_rule_type=nengo.PES()) on2on_ens = nengo.Connection(err_onchip, ens_conn.learning_rule) on2on_neurons = nengo.Connection(err_onchip, neurons_conn.learning_rule) off2on_ens = nengo.Connection(err_offchip, ens_conn.learning_rule) off2on_neurons = nengo.Connection(err_offchip, neurons_conn.learning_rule) networks = SplitNetworks(net, node_neurons=default_node_neurons) networks.move(pre, "chip") networks.move(post, "chip") networks.move(err_onchip, "chip") networks.move(err_offchip, "host") networks.move(ens_conn, "chip") networks.move(neurons_conn, "chip") networks.needs_sender[ens_conn.learning_rule] = "ens_pes_target" networks.needs_sender[neurons_conn.learning_rule] = "neurons_pes_target" split_host_to_learning_rules(networks, networks.original.all_connections) assert on2on_ens not in networks assert on2on_neurons not in networks assert sorted([type(obj).__name__ for obj in networks.adds]) == [ "Connection", "Connection", "HostSendNode", "HostSendNode", ] assert off2on_ens in networks.removes assert "ens_pes_target" in list(networks.host2chip_senders.values()) assert off2on_neurons in networks.removes assert "neurons_pes_target" in list(networks.host2chip_senders.values())
def test_split_chip_to_host(): with nengo.Network() as net: ens_onchip = nengo.Ensemble(10, 1) ens_offchip = nengo.Ensemble(10, 1) node_offchip = nengo.Node(size_in=1) connections = [ nengo.Connection(ens_onchip, ens_offchip), nengo.Connection(ens_onchip, ens_offchip, learning_rule_type=nengo.PES()), nengo.Connection(ens_onchip, node_offchip), nengo.Connection(ens_onchip.neurons, ens_offchip, transform=np.ones((1, 10))), nengo.Connection(ens_onchip.neurons, node_offchip, transform=np.ones((1, 10))), ] connections.append( nengo.Connection(ens_onchip, connections[1].learning_rule)) networks = SplitNetworks(net, node_neurons=default_node_neurons) networks.move(ens_onchip, "chip") networks.move(ens_offchip, "host") networks.move(node_offchip, "host") for conn in connections: split_chip_to_host(networks, conn) for added in networks.adds: if isinstance(added, HostReceiveNode): receive = added elif isinstance(added, nengo.Probe): probe = added else: assert added.post is conn.post receive2post = added assert networks.location(receive) == "host" assert networks.location(receive2post) == "host" assert receive2post.pre is receive assert networks.location(probe) == "chip" assert probe.target is conn.pre or probe.target is conn.pre.ensemble assert probe.synapse is None assert probe in networks.chip2host_params assert probe in networks.chip2host_receivers assert networks.chip2host_receivers[probe] is receive if conn.learning_rule_type is not None: assert conn.learning_rule in networks.needs_sender assert isinstance(networks.needs_sender[conn.learning_rule], PESModulatoryTarget) assert conn in networks.removes networks.adds.clear() # makes next loop iteration easier
def test_split_host_to_chip(): with nengo.Network() as net: ens_offchip = nengo.Ensemble(10, 1) node_offchip = nengo.Node(np.sin) ens_onchip = nengo.Ensemble(10, 1) connections = [ nengo.Connection(ens_offchip, ens_onchip), nengo.Connection(node_offchip, ens_onchip), nengo.Connection(ens_offchip, ens_onchip.neurons, transform=np.ones((10, 1))), nengo.Connection(node_offchip, ens_onchip.neurons, transform=np.ones((10, 1))), ] networks = SplitNetworks(net, node_neurons=default_node_neurons) networks.move(ens_offchip, "host") networks.move(node_offchip, "host") networks.move(ens_onchip, "chip") for conn in connections: split_host_to_chip(networks, conn) for added in networks.adds: if isinstance(added, nengo.Ensemble): ens = added elif isinstance(added, ChipReceiveNode): receive = added elif isinstance(added, HostSendNode): send = added # Otherwise must be connection elif added.pre is conn.pre: pre2ens = added elif added.post is conn.post: receive2post = added else: ensneurons2send = added assert networks.location(ens) == "host" assert isinstance(ens.neuron_type, NIF) assert pre2ens.post is ens assert networks.location(receive) == "chip" assert networks.location(receive2post) == "chip" assert receive2post.pre is receive assert networks.location(send) == "host" assert networks.location(ensneurons2send) == "host" assert ensneurons2send.pre == ens.neurons assert ensneurons2send.post is send assert conn in networks.removes networks.adds.clear() # makes next loop iteration easier
def test_bad_obj_type(): split = SplitNetworks(nengo.Network()) split.adds = {"woops": "host"} with pytest.raises(AssertionError, match="cannot handle type"): split.finalize()
def test_splitnetwork_bad_add_type(): net = nengo.Network() networks = SplitNetworks(net) networks.add(1, "chip") with pytest.raises(AssertionError): networks.finalize()