Function test_apply

tests/python/tirx/test_layout.py:1067–1252

()

def test_apply():
    ################ TileLayout
    def test_tile_layout_0():
        layout = TileLayout(S[(8, 8) : (8, 1)])
        for i, j in itertools.product(range(8), range(8)):
            assert layout.apply(i * 8 + j)["m"] == i * 8 + j * 1
        for i, j in itertools.product(range(8), range(8)):
            assert layout.apply(i, j, shape=(8, 8))["m"] == i * 8 + j * 1
        # # apply can accept coord larger than size
        # for p in range(1024):
        #     outer = p // 64
        #     inner = p % 64
        #     i, j = inner // 8, inner % 8
        #     assert layout.apply(p)["m"] == outer * 64 + i * 8 + j * 1
        with pytest.raises(Exception):
            layout.apply(1, 1, 1)

    test_tile_layout_0()

    def test_tile_layout_1():
        layout = TileLayout(S[(8, 8) : (10, 1)])
        for i, j in itertools.product(range(8), range(8)):
            assert layout.apply(i * 8 + j)["m"] == i * 10 + j * 1
        for i, j in itertools.product(range(8), range(8)):
            assert layout.apply(i, j, shape=(8, 8))["m"] == i * 10 + j * 1

        # # apply can accept coord larger than size
        # for p in range(1024):
        #     outer = p // 64
        #     inner = p % 64
        #     i, j = inner // 8, inner % 8
        #     assert (
        #         layout.apply(
        #             p,
        #         )[0]
        #         == outer * 78 + i * 10 + j * 1
        #     )

    test_tile_layout_1()

    def test_tile_layout_2():
        layout = TileLayout(S[(2, 3, 4, 2, 2) : (1, 2, 12, 6, 48)])

        def f(i0, i1):
            leaf1 = i0 // 3
            leaf2 = i0 % 3
            leaf3 = i1 // 4
            leaf4 = (i1 % 4) // 2
            leaf5 = i1 % 2
            assert (
                layout.apply(i0, i1, shape=(6, 16))["m"]
                == leaf1 * 1 + leaf2 * 2 + leaf3 * 12 + leaf4 * 6 + leaf5 * 48
            )

        for i0, i1 in itertools.product(range(6), range(16)):
            f(i0, i1)
        for i in range(6 * 16):
            f(i // 16, i % 16)