test_lang_container.py 954 Bytes
Newer Older
tqchen committed
1 2 3 4 5 6
import tvm

def test_array():
    a = tvm.convert([1,2,3])
    assert len(a) == 3

7 8 9 10 11 12
def test_array_save_load_json():
    a = tvm.convert([1,2,3])
    json_str = tvm.save_json(a)
    a_loaded = tvm.load_json(json_str)
    assert(a[1].value == 2)

tqchen committed
13
def test_map():
14 15
    a = tvm.var('a')
    b = tvm.var('b')
tqchen committed
16 17 18 19 20
    amap = tvm.convert({a: 2,
                        b: 3})
    assert a in amap
    assert len(amap) == 2
    dd = dict(amap.items())
21 22
    assert a in dd
    assert b in dd
tqchen committed
23 24
    assert a + 1 not in amap

25
def test_map_save_load_json():
26 27
    a = tvm.var('a')
    b = tvm.var('b')
28 29 30 31 32 33 34 35 36
    amap = tvm.convert({a: 2,
                        b: 3})
    json_str = tvm.save_json(amap)
    amap = tvm.load_json(json_str)
    assert len(amap) == 2
    dd = {kv[0].name : kv[1].value for kv in amap.items()}
    assert(dd == {"a": 2, "b": 3})


tqchen committed
37 38 39
if __name__ == "__main__":
    test_array()
    test_map()
40 41
    test_array_save_load_json()
    test_map_save_load_json()