Function test_top

dask/array/tests/test_array_core.py:130–218

()

def test_top():
    t = Task("inc", inc, TaskRef("A"))
    assert _make_blockwise_graph(
        t, "z", "ij", "x", "ij", numblocks={"x": (2, 2)}, keys=("A",)
    ) == {
        ("z", 0, 0): Task(("z", 0, 0), inc, TaskRef(("x", 0, 0))),
        ("z", 0, 1): Task(("z", 0, 1), inc, TaskRef(("x", 0, 1))),
        ("z", 1, 0): Task(("z", 1, 0), inc, TaskRef(("x", 1, 0))),
        ("z", 1, 1): Task(("z", 1, 1), inc, TaskRef(("x", 1, 1))),
    }

    keys = list(map(str, range(2)))
    t = Task("add", add, *(TaskRef(k) for k in keys))
    assert _make_blockwise_graph(
        t,
        "z",
        "ij",
        "x",
        "ij",
        "y",
        "ij",
        numblocks={"x": (2, 2), "y": (2, 2)},
        keys=keys,
    ) == {
        ("z", 0, 0): Task(("z", 0, 0), add, TaskRef(("x", 0, 0)), TaskRef(("y", 0, 0))),
        ("z", 0, 1): Task(("z", 0, 1), add, TaskRef(("x", 0, 1)), TaskRef(("y", 0, 1))),
        ("z", 1, 0): Task(("z", 1, 0), add, TaskRef(("x", 1, 0)), TaskRef(("y", 1, 0))),
        ("z", 1, 1): Task(("z", 1, 1), add, TaskRef(("x", 1, 1)), TaskRef(("y", 1, 1))),
    }
    keys = list(map(str, range(2)))
    t = Task(
        "dotmany",
        dotmany,
        TaskRef(keys[0]),
        TaskRef(keys[1]),
    )
    out = _make_blockwise_graph(
        t,
        "z",
        "ik",
        "x",
        "ij",
        "y",
        "jk",
        numblocks={"x": (2, 2), "y": (2, 2)},
        keys=keys,
    )
    expected = {
        ("z", 0, 0): Task(
            ("z", 0, 0),
            dotmany,
            List(TaskRef(("x", 0, 0)), TaskRef(("x", 0, 1))),
            List(TaskRef(("y", 0, 0)), TaskRef(("y", 1, 0))),
        ),
        ("z", 0, 1): Task(
            ("z", 0, 1),
            dotmany,
            List(TaskRef(("x", 0, 0)), TaskRef(("x", 0, 1))),