def test_optimize_with_getitem_fusion(): dsk = {'a': 'some-array', 'b': (getarray, 'a', (slice(10, 20), slice(100, 200))), 'c': (getarray, 'b', (5, slice(50, 60)))} result = optimize(dsk, ['c']) expected = {'c': (getarray, 'some-array', (15, slice(150, 160)))} assert result == expected
def test_optimize_with_getitem_fusion(): dsk = {'a': 'some-array', 'b': (getter, 'a', (slice(10, 20), slice(100, 200))), 'c': (getter, 'b', (5, slice(50, 60)))} result = optimize(dsk, ['c']) expected_task = (getter, 'some-array', (15, slice(150, 160))) assert any(v == expected_task for v in result.values()) assert len(result) < len(dsk)
def test_fuse_slices_with_alias(): dsk = {'x': np.arange(16).reshape((4, 4)), ('dx', 0, 0): (getter, 'x', (slice(0, 4), slice(0, 4))), ('alias', 0, 0): ('dx', 0, 0), ('dx2', 0): (getitem, ('alias', 0, 0), (slice(None), 0))} keys = [('dx2', 0)] dsk2 = optimize(dsk, keys) assert len(dsk2) == 3 fused_key = set(dsk2).difference(['x', ('dx2', 0)]).pop() assert dsk2[fused_key] == (getter, 'x', (slice(0, 4), 0))
def test_optimize_with_getitem_fusion(): dsk = { "a": "some-array", "b": (getter, "a", (slice(10, 20), slice(100, 200))), "c": (getter, "b", (5, slice(50, 60))), } result = optimize(dsk, ["c"]) expected_task = (getter, "some-array", (15, slice(150, 160))) assert any(v == expected_task for v in result.values()) assert len(result) < len(dsk)
def test_fuse_slices_with_alias(): dsk = { "x": np.arange(16).reshape((4, 4)), ("dx", 0, 0): (getter, "x", (slice(0, 4), slice(0, 4))), ("alias", 0, 0): ("dx", 0, 0), ("dx2", 0): (getitem, ("alias", 0, 0), (slice(None), 0)), } keys = [("dx2", 0)] dsk2 = optimize(dsk, keys) assert len(dsk2) == 3 fused_key = set(dsk2).difference(["x", ("dx2", 0)]).pop() assert dsk2[fused_key] == (getter, "x", (slice(0, 4), 0))
def test_fuse_slices_with_alias(getter, getitem): dsk = { "x": np.arange(16).reshape((4, 4)), ("dx", 0, 0): (getter, "x", (slice(0, 4), slice(0, 4))), ("alias", 0, 0): ("dx", 0, 0), ("dx2", 0): (getitem, ("alias", 0, 0), (slice(None), 0)), } keys = [("dx2", 0)] dsk2 = optimize(dsk, keys) assert len(dsk2) == 3 fused_key = (dsk2.keys() - {"x", ("dx2", 0)}).pop() assert _check_get_task_eq(dsk2[fused_key], (getter, "x", (slice(0, 4), 0)))
def test_disable_lowlevel_fusion(): """Check that by disabling fusion, the HLG survives through optimizations""" with dask.config.set({"optimization.fuse.active": False}): y = da.ones(3, chunks=(3, ), dtype="int") optimize = y.__dask_optimize__ dsk1 = y.__dask_graph__() dsk2 = optimize(dsk1, y.__dask_keys__()) assert isinstance(dsk1, HighLevelGraph) assert isinstance(dsk2, HighLevelGraph) assert dsk1 == dsk2 y = y.persist() assert isinstance(y.__dask_graph__(), HighLevelGraph) assert_eq(y, [1] * 3)
def test_dont_fuse_different_slices(): x = da.random.random(size=(10, 10), chunks=(10, 1)) y = x.rechunk((1, 10)) dsk = optimize(y.dask, y._keys()) assert len(dsk) > 100