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)
