graph: FlowGraph
model: hidet.runtime.CompiledModel = graph.build()
outputs = model(inputs)
model.save("model.hidet")
model_loaded = hidet.load_model("model.hidet")

def test_load_save(device: str):
    # construct graph
    x = hidet.symbol([2, 3], device=device)
    w1 = hidet.randn([3, 4], device=device)
    w2 = hidet.randn([4, 5], device=device)
    y = hidet.ops.matmul(hidet.ops.matmul(x, w1), w2)

    # get computation graph
    graph = hidet.trace_from(y)

    # optimize the graph
    graph = hidet.graph.optimize(graph)

    # build the graph
    model = graph.build()

    # save the model
    model.save('./model.hidet')

    # load the model
    loaded_model = hidet.load_model('./model.hidet')

    # compare the results
    xx = hidet.randn([2, 3], device=device)
    y1 = graph(xx)
    y2 = model(xx)
    y3 = loaded_model(xx)

    numpy.testing.assert_allclose(y1.cpu().numpy(), y2.cpu().numpy())
    numpy.testing.assert_allclose(y1.cpu().numpy(), y3.cpu().numpy())

import hidet


def test_args():
    from hidet.lang import attrs, meta, printf, int32

    with hidet.script_module() as script_module:

        @hidet.script
        def launch(args: meta.types([int, bool, float, int32]), second_args: int, thrid_args: meta.types([int32])):
            attrs.func_kind = 'public'

            printf("%d\n", args[0])
            printf("%d\n", args[1])
            printf("%f\n", args[2])
            printf("%d\n", args[3])
            printf("%d\n", second_args)
            printf("%d\n", thrid_args[0])

    module = script_module.build()
    module(1, True, 0.1, 2, 3, 4)
    print(module.source(True))