Relay中的函数的作用类似于其他编程语言中的过程或函数,并用于推广命名子图的概念。函数是Relay中的第一类,这意味着它们就像变量、常量和元组一样是表达式。 另外,Relay中的函数是高阶的,这意味着可以将函数作为参数传递给函数或由函数返回,因为函数表达式求值闭包,闭包是张量和元组之类的值 。函数节点的具体定义可以和文档可以参考https://docs.tvm.ai/api/python/relay/expr.html#tvm.relay.expr.Function。
语法:
定义至少由关键字fn,空参数集和花括号所包含的主体表达式(Expr)组成。
fn() { body }
定义也可以包含任意数目的参数。例如一个涉及add操作符的简单函数。
fn(%x, %y) { add(%x, %y) }
注意,在函数体中,参数是局部变量,就像let表达式中绑定的那些变量一样。
还可以在函数上显式注释类型。例如,我们可以限制上面的函数只在某些类型上工作。函数仅接受Tensor [(10,10),float32]类型的参数,并返回Tensor [(10,10),float32]类型的值。 函数参数只是一个局部变量(LocalVar),可以选择使用类型注释,写为%x:T。
fn(%x : Tensor[(10, 10), float32], %y : Tensor[(10, 10), float32])
-> Tensor[(10, 10), float32] {
add(%x, %y)
}
当类型信息被省略时,Relay尝试为用户推断最通用的类型。这个属性称为一般化:对于没有显式注释的定义,Relay尝试将最一般的类型分配给参数,并根据函数体和调用站点返回类型。
递归函数表达式可以使用let绑定定义,如下所示:
let %fact = fn(%x : Tensor[(10, 10), float32]) -> Tensor[(10, 10), float32] {
if (%x == Constant(0, (10, 10), float32)) {
Constant(1, (10, 10), float32)
} else {
%x * %fact(%x - Constant(1, (10, 10), float32))
}
};
%fact(Constant(10, (10, 10), float32))
闭包:
函数表达式的计算结果为闭包。 闭包是表示为一对局部环境(存储在函数体范围之外定义的所有变量的值)和函数本身的值。例如,在下面的例子中,最终的结果将是一个0值的张量,因为%f的闭包存储了%x在定义%f的指针处的值。
多态与类型关系:
注意:文本格式尚不支持类型参数语法。
还可以为函数提供一组类型参数,这些参数可以代替调用站点上的特定类型。 具有类型参数的函数是类型多态的; 它们的返回类型或将接受的参数类型可以根据调用站点提供的类型参数而有所不同。类型参数是按种类分类的,并且只能出现在类型签名合适的部分(例如,类型为Shape的类型参数只能出现在张量类型中期望出现形状的位置); 有关完整的讨论,请参见https://docs.tvm.ai/langref/relay_type.html#type-parameter。例如,可以为任何Relay类型定义一个多态标识函数,如下所示:
fn<t : Type>(%x : t) -> t {
%x
}
以下定义也是多态的,但将其参数限制为张量类型:
fn<s : Shape, bt : BaseType>(%x : Tensor[s, bt]) {
%x
}
请注意,返回类型将被省略,并将被推断。
注意:文本格式尚不支持“ where”语法。
函数也可能受一个或多个类型关系的约束,例如:
fn(%x, %y) where Broadcast { add(%x, %y) }
在上面的定义中,%x和%y的类型以及返回类型均受Broadcast关系的约束,这意味着所有这三个必须为张量,并且它们的形状遵循元素级广播关系。 与运算符一样,关系的定义对Relay不是透明的,而是在外部以C ++或Python实现。与广播一样,关系用于表达对类型(尤其是张量形状)的复杂约束。 所有函数关系必须在所有调用位置都存在; 因此,将类型检查视为约束解决问题。 有关类型关系及其实现的更多详细信息,请参见https://docs.tvm.ai/langref/relay_type.html#type-relation。