Error splitting probe across blocks

[studywolf]

import tensorflow as tf

import nengo
import nengo_dl
import nengo_loihi

inp = tf.keras.Input(shape=(28, 28, 1))
conv0 = tf.keras.layers.Conv2D(
    filters=3,  
    strides=1,
    kernel_size=1,
    activation=tf.nn.relu,
    use_bias=False,
)(inp)

# convolutional layers
conv1 = tf.keras.layers.Conv2D(
    filters=32,
    kernel_size=3,
    activation=tf.nn.relu,
    use_bias=False,
)(conv0)

model = tf.keras.Model(inputs=inp, outputs=conv1)  

# convert the keras model to a nengo network
nengo_converter = nengo_dl.Converter(
    model,
    swap_activations={tf.nn.relu: nengo_loihi.neurons.LoihiSpikingRectifiedLinear()},
)

# add functionality to input node, to send in images over time
with nengo_converter.net as net:
    nengo_loihi.add_params(net)

    net.config[nengo_converter.layers[conv0].ensemble].on_chip = False
    net.config[nengo_converter.layers[conv1].ensemble].block_shape = (
        nengo_loihi.BlockShape((10, 10, 1), (28, 28, 1)))

# build Nengo Loihi Simulator and run network
with nengo_loihi.Simulator(nengo_converter.net) as sim:
    sim.run(1)

gives

ids:  (784,)
old_comp:  (21632,)
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-22-9aba82875b45> in <module>
     40 
     41 # build Nengo Loihi Simulator and run network
---> 42 with nengo_loihi.Simulator(nengo_converter.net) as sim:
     43     sim.run(1)

~/Dropbox/code/nengo-loihi/nengo_loihi/simulator.py in __init__(self, network, dt, seed, model, precompute, target, progress_bar, remove_passthrough, hardware_options)
    137             precompute=precompute,
    138             remove_passthrough=remove_passthrough,
--> 139             discretize=target != "simreal",
    140         )
    141 

~/Dropbox/code/nengo-loihi/nengo_loihi/builder/builder.py in build(self, obj, *args, **kwargs)
    220             add_params(obj)
    221 
--> 222         built = model.builder.build(model, obj, *args, **kwargs)
    223         if self.build_callback is not None:
    224             self.build_callback(obj)

~/Dropbox/code/nengo/nengo/builder/builder.py in build(cls, model, obj, *args, **kwargs)
    240             raise BuildError("Cannot build object of type %r" % type(obj).__name__)
    241 
--> 242         return cls.builders[obj_cls](model, obj, *args, **kwargs)
    243 
    244     @classmethod

~/Dropbox/code/nengo-loihi/nengo_loihi/builder/network.py in build_network(model, network, precompute, remove_passthrough, discretize)
     42 
     43         # Split blocks into blocks that will fit on cores
---> 44         split_model(model)
     45 
     46         if discretize:

~/Dropbox/code/nengo-loihi/nengo_loihi/builder/split_blocks.py in split_model(model)
    129     for probe in model.probes:
    130         print(probe.target)
--> 131         split_probe(probe, block_map, synapse_map)
    132 
    133     new_blocks = [block for group in block_map.values() for block in group]

~/Dropbox/code/nengo-loihi/nengo_loihi/builder/split_blocks.py in split_probe(probe, block_map, synapse_map)
    208     print('ids: ', ids.shape)
    209     print('old_comp: ', old_comp_ids.shape)
--> 210     assert ids.shape == old_comp_ids.shape
    211     assert np.array_equal(np.unique(ids), old_comp_ids)
    212 

AssertionError: 

[hunse]

It's because the ensemble shape (second parameter) of your block shape doesn't match the number of neurons in the ensemble. We should have a check for this to let you know; I'll add something.

[hunse]

You can pass in the convolution transform, and BlockShape will get the output shape from it. https://www.nengo.ai/nengo-loihi/api.html#nengo_loihi.BlockShape

[studywolf]

great, thanks!