test_lang_buffer.py 353 Bytes
Newer Older
1 2 3
import tvm

def test_buffer():
4 5 6 7 8
    m = tvm.var('m')
    n = tvm.var('n')
    l = tvm.var('l')
    Ab = tvm.decl_buffer((m, n), tvm.float32)
    Bb = tvm.decl_buffer((n, l), tvm.float32)
9

10
    assert isinstance(Ab, tvm.schedule.Buffer)
11 12 13 14 15 16
    assert Ab.dtype == tvm.float32
    assert tuple(Ab.shape) == (m, n)


if __name__ == "__main__":
    test_buffer()