def test_block_shape_errors():
with pytest.raises(ValidationError, match="[Mm]ust be a tuple"):
BlockShape([5], [15])
with pytest.raises(ValidationError, match="[Mm]ust be an int"):
BlockShape((5, ), (15.0, ))
with pytest.raises(ValidationError, match="[Mm]ust be the same length"):
BlockShape((2, 2), (8, ))
def test_block_shape_3d():
block_shape = BlockShape((2, 3, 2), (4, 6, 7))
assert block_shape.shape == (2, 3, 2)
assert block_shape.ensemble_shape == (4, 6, 7)
assert block_shape.n_splits == 2 * 2 * 4
assert block_shape.block_size == 2 * 3 * 2
assert block_shape.ensemble_size == 4 * 6 * 7
assert list(block_shape.zip_dimensions()) == [(4, 2), (6, 3), (7, 2)]
def test_block_shape_1d():
block_shape = BlockShape((5,), (15,))
assert block_shape.shape == (5,)
assert block_shape.ensemble_shape == (15,)
assert block_shape.n_splits == 3
assert block_shape.block_size == 5
assert block_shape.ensemble_size == 15
assert list(block_shape.zip_dimensions()) == [(15, 5)]
def test_block_shape_3d():
block_shape = BlockShape((2, 3, 2), (4, 6, 7))
assert block_shape.shape == (2, 3, 2)
assert block_shape.ensemble_shape == (4, 6, 7)
assert block_shape.n_splits == 2 * 2 * 4
assert block_shape.block_size == 2 * 3 * 2
assert block_shape.ensemble_size == 4 * 6 * 7
assert list(block_shape.zip_dimensions()) == [(4, 2), (6, 3), (7, 2)]
with pytest.warns(UserWarning, match="uses the same number of blocks as"):
block_shape = BlockShape((2, 3, 5), (4, 6, 7))
assert block_shape.shape == (2, 3, 4)
assert block_shape.n_splits == 2 * 2 * 2
def test_block_shape_1d():
block_shape = BlockShape((5, ), (15, ))
assert block_shape.shape == (5, )
assert block_shape.ensemble_shape == (15, )
assert block_shape.n_splits == 3
assert block_shape.block_size == 5
assert block_shape.ensemble_size == 15
assert list(block_shape.zip_dimensions()) == [(15, 5)]
with pytest.warns(UserWarning, match="uses the same number of blocks as"):
block_shape = BlockShape((4, ), (9, ))
assert block_shape.shape == (3, )
assert block_shape.n_splits == 3
def test_block_shape_errors():
with pytest.raises(ValidationError, match="[Mm]ust be a tuple"):
BlockShape([5], [15])
with pytest.raises(ValidationError, match="[Mm]ust be an int"):
BlockShape((5,), (15.0,))
with pytest.raises(ValidationError, match="[Mm]ust be the same length"):
BlockShape((2, 2), (8,))
with nengo.Network() as net:
add_params(net)
a = nengo.Ensemble(10, 1)
with pytest.raises(ValidationError, match="Block shape ensemble size"):
net.config[a].block_shape = BlockShape((3, 2), (6, 2))
def split_block(old_block, block_shapes, validate=ValidationLevel.MINIMAL):
"""Break a block apart into smaller blocks, each able to fit on one core"""
n_compartments = old_block.compartment.n_compartments
n_in_axons = sum(synapse.n_axons for synapse in old_block.synapses)
n_out_axons = sum(axon.axon_slots() for axon in old_block.axons)
synapse_bits = sum(synapse.bits() for synapse in old_block.synapses)
if block_shapes.get(old_block, None) is None:
# break block sequentially
# TODO: account for compartments having different numbers of synapses/axons/etc.
# Splitting into blocks where each block has the same number of compartments
# could leave blocks that have more synapses or axons than allowed. But this
# is rare, and users can work around it by specifying the split shape manually
n_split = max((
ceil_div(n_compartments, MAX_COMPARTMENTS),
ceil_div(n_in_axons, MAX_IN_AXONS),
ceil_div(n_out_axons, MAX_OUT_AXONS),
ceil_div(synapse_bits, MAX_SYNAPSE_BITS),
))
block_shapes[old_block] = BlockShape(
(ceil_div(n_compartments, n_split), ), (n_compartments, ))
old_block_shape = block_shapes[old_block]
assert old_block_shape.ensemble_size == old_block.n_neurons
# find compartment indices for each new block
new_block_inds = []
ranges = [range(0, n, i) for n, i in old_block_shape.zip_dimensions()]
full_inds = np.arange(old_block_shape.ensemble_size).reshape(
old_block_shape.ensemble_shape)
for inds0 in itertools.product(*ranges):
inds1 = np.minimum(inds0 + old_block_shape._shape,
old_block_shape._ens_shape)
indslice = tuple(slice(i0, i1) for i0, i1 in zip(inds0, inds1))
inds = full_inds[indslice]
new_block_inds.append(IndicesList(inds.flat))
assert len(new_block_inds) > 0
if len(new_block_inds) == 1:
# if block can fit on one core, just return the current block
if validate >= ValidationLevel.MINIMAL:
assert new_block_inds[0].set == set(range(n_compartments))
new_blocks = [old_block]
return OrderedDict(zip(new_blocks, new_block_inds))
# break apart block
new_blocks = []
for k, inds in enumerate(new_block_inds):
n_neurons = len(inds)
new_block = LoihiBlock(n_neurons)
if old_block.label is not None:
ind_array = np.array(list(inds))
d = np.diff(ind_array)
indstr = ("%d:%d:%d" % (ind_array[0], ind_array[-1] + 1, d[0])
if len(d) > 0 and np.all(d[0] == d) else "%d:%d" %
(ind_array[0], ind_array[0] + 1)
if len(ind_array) == 1 else str(k))
new_block.label = "%s[%s]" % (old_block.label, indstr)
for attr in (
"decay_u",
"decay_v",
"refract_delay",
"vth",
"bias",
"enable_noise",
):
# copy whole array to ensure that we maintain dtype
setattr(
new_block.compartment,
attr,
getattr(old_block.compartment, attr)[list(inds)].copy(),
)
for attr in (
"tau_s",
"scale_u",
"scale_v",
"vmin",
"vmax",
"noise_offset",
"noise_exp",
"noise_at_membrane",
):
setattr(new_block.compartment, attr,
getattr(old_block.compartment, attr))
new_blocks.append(new_block)
logger.info(
"Split block (%d) into (%s)",
n_compartments,
", ".join(
str(new_block.compartment.n_compartments)
for new_block in new_blocks),
)
return OrderedDict(zip(new_blocks, new_block_inds))