Relay 函数

Relay 中的函数的作用类似于其他编程语言中的过程或函数，并用于推广命名子图的概念。

可以直接定义函数：

from tvm import relay

a, b = [relay.var(name) for name in "ab"]
add_op = a + b
add_func = relay.Function([a, b], add_op)

此函数的文本形式：

add_func

fn (%a, %b) {
  add(%a, %b)
}

也可以使用 Python 函数回调的形式：

def add(a, b):
    add_op = a + b
    return relay.Function([a, b], add_op)

add

<function __main__.add(a, b)>

想要使用，需要：

a, b = [relay.var(name) for name in "ab"]
add_func = add(a, b)
add_func

fn (%a, %b) {
  add(%a, %b)
}

也可以添加变量注解：

type_annotation = relay.TensorType(shape=(5, 5),
                                   dtype="float32")

def add(a, b):
    add_op = a + b
    return relay.Function([a, b],
                          add_op,
                          ret_type=type_annotation,
                          type_params=None)


a, b = [relay.var(name, type_annotation) for name in "ab"]
add_func = add(a, b)
add_func

fn (%a: Tensor[(5, 5), float32], %b: Tensor[(5, 5), float32]) -> Tensor[(5, 5), float32] {
  add(%a, %b)
}