def build_connection(model, conn): pre_onchip = model.split.on_chip(base_obj(conn.pre)) if isinstance(conn.post_obj, LearningRule): assert not pre_onchip return build_host_to_learning_rule(model, conn) post_onchip = model.split.on_chip(base_obj(conn.post)) if pre_onchip and post_onchip: build_chip_connection(model, conn) elif not pre_onchip and post_onchip: if isinstance(conn.pre_obj, Neurons): build_host_neurons_to_chip(model, conn) else: build_host_to_chip(model, conn) elif pre_onchip and not post_onchip: build_chip_to_host(model, conn) else: assert not pre_onchip and not post_onchip host = model.host_model(base_obj(conn.pre)) assert host is model.host_model(base_obj(conn.post)) _inherit_seed(host, conn, model, conn) host.build(conn)
def build_host_to_learning_rule(model, conn): if not is_transform_type(conn.transform, ("Dense", "NoTransform")): # TODO: What needs to be done to support this? It looks like it should just work raise BuildError( f"{conn}: nengo-loihi does not yet support " f"'{type(conn.transform).__name__}' transforms on host to chip " "learning rule connections" ) dim = conn.size_out host = model.host_model(base_obj(conn.pre)) logger.debug("Creating HostSendNode for %s", conn) send = HostSendNode( dim, label=None if conn.label is None else "%s_send" % conn.label, add_to_container=False, ) host.build(send) pre2send = Connection( conn.pre, send, function=conn.function, solver=conn.solver, eval_points=conn.eval_points, scale_eval_points=conn.scale_eval_points, synapse=conn.synapse, transform=conn.transform, label=conn.label, add_to_container=False, ) model.host2chip_pes_senders[send] = model.needs_sender[conn.post_obj] _inherit_seed(host, pre2send, model, conn) host.build(pre2send)
def build_host_to_learning_rule(model, conn): if (nengo_transforms is not None and not isinstance(conn.transform, nengo_transforms.Dense)): raise BuildError("nengo-loihi does not yet support %r transforms " "on host to chip learning rule connections" % (type(conn.transform).__name__, )) dim = conn.size_out host = model.host_model(base_obj(conn.pre)) logger.debug("Creating HostSendNode for %s", conn) send = HostSendNode( dim, label=None if conn.label is None else "%s_send" % conn.label, add_to_container=False, ) host.build(send) pre2send = Connection( conn.pre, send, function=conn.function, solver=conn.solver, eval_points=conn.eval_points, scale_eval_points=conn.scale_eval_points, synapse=conn.synapse, transform=conn.transform, label=conn.label, add_to_container=False, ) pes_target = model.needs_sender[conn.post_obj] model.host2chip_senders[send] = pes_target _inherit_seed(host, pre2send, model, conn) host.build(pre2send)
def build_chip_to_host(model, conn): if not is_transform_type(conn.transform, ("Dense", "NoTransform")): raise BuildError( f"{conn}: nengo-loihi does not yet support " f"'{type(conn.transform).__name__}' transforms on chip to host connections" ) rng = np.random.RandomState(model.seeds[conn]) dim = conn.size_out host = model.host_model(base_obj(conn.post)) logger.debug("Creating HostReceiveNode for %s", conn) receive = HostReceiveNode( dim, label=None if conn.label is None else "%s_receive" % conn.label, add_to_container=False, ) host.build(receive) receive2post = Connection( receive, conn.post, synapse=conn.synapse, label=None if conn.label is None else "%s_host" % conn.label, add_to_container=False, ) _inherit_seed(host, receive2post, model, conn) host.build(receive2post) logger.debug("Creating Probe for %s", conn) transform = sample_transform(conn, rng=rng) probe = NengoProbe( conn.pre, synapse=None, solver=conn.solver, add_to_container=False ) model.chip2host_params[probe] = dict( learning_rule_type=conn.learning_rule_type, function=conn.function, eval_points=conn.eval_points, scale_eval_points=conn.scale_eval_points, transform=transform, label=None if conn.label is None else "%s_probe" % conn.label, ) model.chip2host_receivers[probe] = receive _inherit_seed(model, probe, model, conn) model.builder.build(model, probe) if conn.learning_rule_type is not None: if not isinstance(conn.pre_obj, Ensemble): raise NotImplementedError( "Learning rule presynaptic object must be an Ensemble " "(got %r)" % type(conn.pre_obj).__name__ ) model.needs_sender[conn.learning_rule] = PESModulatoryTarget(probe)
def build_host_to_chip(model, conn): rng = np.random.RandomState(model.seeds[conn]) host = model.host_model(base_obj(conn.pre)) if is_transform_type(conn.transform, ("Convolution", "ConvolutionTranspose")): raise BuildError( f"{conn}: Conv2D transforms not supported for off-chip to " "on-chip connections where `pre` is not a Neurons object." ) elif not is_transform_type(conn.transform, ("Dense", "NoTransform")): raise BuildError( f"{conn}: nengo-loihi does not yet support " f"'{type(conn.transform).__name__}' transforms on host to chip connections" ) # Scale the input spikes based on the radius of the target ensemble weights = sample_transform(conn, rng=rng) if isinstance(conn.post_obj, Ensemble): weights = weights / conn.post_obj.radius if is_transform_type(conn.transform, "NoTransform"): transform = weights # weights are 1 / (post ensemble radius), if applicable else: # copy the Transform information, setting `init` to the sampled weights transform = copy.copy(conn.transform) type(transform).init.data[transform] = weights if isinstance(conn.post_obj, Neurons): # we don't have encoders, and the transform could have large output, # so do it on the chip host_transform = 1.0 chip_transform = transform dim = conn.size_mid else: # we have encoders on the chip, so do the transform off-chip host_transform = transform chip_transform = 1.0 dim = conn.size_out logger.debug("Creating ChipReceiveNode for %s", conn) receive = ChipReceiveNode( dim * 2, size_out=dim, label=None if conn.label is None else "%s_node" % conn.label, add_to_container=False, ) model.builder.build(model, receive) receive2post = Connection( receive, conn.post, transform=chip_transform, synapse=model.decode_tau, label=None if conn.label is None else "%s_chip" % conn.label, add_to_container=False, ) _inherit_seed(model, receive2post, model, conn) _inherit_config(model, receive2post, model, conn) build_chip_connection(model, receive2post) logger.debug("Creating DecodeNeuron ensemble for %s", conn) ens = model.node_neurons.get_ensemble(dim, add_to_container=False) ens.label = None if conn.label is None else "%s_ens" % conn.label _inherit_seed(host, ens, model, conn) host.build(ens) model.connection_decode_neurons[conn] = ens pre2ens = Connection( conn.pre, ens, function=conn.function, solver=conn.solver, eval_points=conn.eval_points, scale_eval_points=conn.scale_eval_points, synapse=conn.synapse, transform=host_transform, label=None if conn.label is None else "%s_enc" % conn.label, add_to_container=False, ) _inherit_seed(host, pre2ens, model, conn) host.build(pre2ens) logger.debug("Creating HostSendNode for %s", conn) send = HostSendNode( dim * 2, label=None if conn.label is None else "%s_send" % conn.label, add_to_container=False, ) host.build(send) ensneurons2send = Connection( ens.neurons, send, synapse=None, label=None if conn.label is None else "%s_host" % conn.label, add_to_container=False, ) _inherit_seed(host, ensneurons2send, model, conn) model.host2chip_senders[send] = receive host.build(ensneurons2send)