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_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():
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
