Get to the bottom of the new dask version error

[daffidwilde]

Workaround in #144, but this error is raised with the lateat Dask version (2012.12.0):

(edo-dev) edo(fix-failing-action): python -m pytest -v tests/test_optimiser.py::test_get_pop_history
======================================================== test session starts ========================================================
platform darwin -- Python 3.9.1, pytest-6.2.1, py-1.10.0, pluggy-0.12.0 -- /Users/henrywilde/opt/anaconda3/envs/edo-dev/bin/python
cachedir: .pytest_cache
hypothesis profile 'default' -> database=DirectoryBasedExampleDatabase('/Users/henrywilde/src/edo/.hypothesis/examples')
rootdir: /Users/henrywilde/src/edo
plugins: hypothesis-5.20.3, cov-2.10.1
collected 1 item

tests/test_optimiser.py::test_get_pop_history FAILED                                                                          [100%]

============================================================= FAILURES ==============================================================
_______________________________________________________ test_get_pop_history ________________________________________________________

>   ???

tests/test_optimiser.py:509:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tests/test_optimiser.py:560: in test_get_pop_history
    assert np.allclose(
../../opt/anaconda3/envs/edo-dev/lib/python3.9/site-packages/dask/base.py:279: in compute
    (result,) = compute(self, traverse=False, **kwargs)
../../opt/anaconda3/envs/edo-dev/lib/python3.9/site-packages/dask/base.py:561: in compute
    dsk = collections_to_dsk(collections, optimize_graph, **kwargs)
../../opt/anaconda3/envs/edo-dev/lib/python3.9/site-packages/dask/base.py:332: in collections_to_dsk
    _opt = opt(dsk, keys, **kwargs)
../../opt/anaconda3/envs/edo-dev/lib/python3.9/site-packages/dask/array/optimization.py:48: in optimize
    dsk = dsk.cull(set(keys))
../../opt/anaconda3/envs/edo-dev/lib/python3.9/site-packages/dask/highlevelgraph.py:627: in cull
    all_ext_keys = self.get_all_external_keys()
../../opt/anaconda3/envs/edo-dev/lib/python3.9/site-packages/dask/highlevelgraph.py:518: in get_all_external_keys
    self._all_external_keys.update(layer.get_output_keys())
../../opt/anaconda3/envs/edo-dev/lib/python3.9/site-packages/dask/blockwise.py:285: in get_output_keys
    *[range(self.dims[i]) for i in self.output_indices]
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

.0 = <tuple_iterator object at 0x7ff03f7bb610>

>           *[range(self.dims[i]) for i in self.output_indices]
        )
    }
E   KeyError: 'j'

../../opt/anaconda3/envs/edo-dev/lib/python3.9/site-packages/dask/blockwise.py:285: KeyError
------------------------------------------------------------ Hypothesis -------------------------------------------------------------
Falsifying example: test_get_pop_history(
    size=2,
    row_limits=[1, 1],
    col_limits=[1, 1],
    distributions=[edo.distributions.discrete.Bernoulli,
     edo.distributions.continuous.Uniform],
    weights=(0.01, 0.99),
    max_iter=1,
    best_prop=0.5,
    lucky_prop=0.0,
    crossover_prob=0.0,
    mutation_prob=0.0,
    shrinkage=0.0,
    maximise=False,
)
====================================================== short test summary info ======================================================
FAILED tests/test_optimiser.py::test_get_pop_history - KeyError: 'j'
========================================================= 1 failed in 1.23s =========================================================