Relay let 绑定

Relay let 绑定#

其实上述介绍的模块绑定属于 Relay graph 绑定,对应于计算图。

let 绑定是不可变的局部变量绑定,允许用户将表达式绑定到名称。

  • let 绑定包含局部变量、可选类型注解、值和可以引用绑定标识符的 body 表达式。如果省略了绑定变量上的类型注释,Relay 将尝试推断该变量允许的最通用类型。

  • let 表达式中的绑定变量只作用在其 body 作用域内,除非该变量定义了函数表达式。当 let 表达式创建函数时,该变量的值也在范围内,以允许递归定义函数(请参阅前一小节)。

  • let 绑定的值是计算它所依赖的绑定后的最后一个表达式的值。

from tvm import relay

x = relay.var("x")
sb = relay.ScopeBuilder()
v1 = sb.let("v1", relay.log(x))
v2 = sb.let("v2", v1 + v1)
sb.ret(v2)
f = relay.Function([x], sb.get())

f
fn (%x) {
  let %v1 = log(%x);
  let %v2 = add(%v1, %v1);
  %v2
}

也可以定义 if-else 语句:

sb = relay.ScopeBuilder()
cond = relay.var("cond", 'bool')
x = relay.var("x")
y = relay.var("y")

with sb.if_scope(cond):
    one = relay.const(1, "float32")
    t1 = sb.let("t1", relay.add(x, one))
    sb.ret(t1)
with sb.else_scope():
    sb.ret(y)
f = relay.Function([x, y, cond], sb.get())
f
fn (%x, %y, %cond: bool) {
  if (%cond) {
    let %t1 = add(%x, 1f);
    %t1
  } else {
    %y
  }
}