mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Lei Wang
9bd62f5b5f
* layout: add Layout.expand and improve repeat errors * revert unexpected changes
50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
import pytest
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
from tilelang.layout import Layout
|
|
|
|
tilelang.testing.set_random_seed()
|
|
|
|
|
|
def test_layout_expand_linear_scalar_output():
|
|
base = Layout([16, 128], lambda j, k: j * 128 + k)
|
|
expanded = base.expand([4])
|
|
|
|
expected = Layout([4, 16, 128], lambda i, j, k: [i, j * 128 + k])
|
|
assert expanded.is_equal(expected)
|
|
|
|
assert list(expanded.get_input_shape()) == [4, 16, 128]
|
|
assert list(expanded.get_output_shape()) == [4, 2048]
|
|
|
|
|
|
def test_layout_expand_multi_output_identity():
|
|
base = Layout([16, 128], lambda j, k: [j, k])
|
|
expanded = base.expand(4)
|
|
|
|
expected = Layout([4, 16, 128], lambda i, j, k: [i, j, k])
|
|
assert expanded.is_equal(expected)
|
|
|
|
assert list(expanded.get_output_shape()) == [4, 16, 128]
|
|
|
|
|
|
def test_layout_expand_noop_empty():
|
|
base = Layout([16, 128], lambda j, k: j * 128 + k)
|
|
assert base.expand([]) is base
|
|
assert base.expand(()) is base
|
|
|
|
|
|
def test_layout_expand_invalid_args():
|
|
base = Layout([16, 128], lambda j, k: j * 128 + k)
|
|
|
|
with pytest.raises(ValueError):
|
|
_ = base.expand(0)
|
|
with pytest.raises(ValueError):
|
|
_ = base.expand([-1])
|
|
with pytest.raises(TypeError):
|
|
_ = base.expand("4")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|