gistfile1.txt

import tvm
import tvm.relay as relay
'''
fn (%cond: bool) -> float32 {
  let %x = fn () -> float32 {
    let %x1 = 0f /* ty=float32 */
    %x1
  }
  let %x2 = ref(%x)
  let %x3 = if (%cond) {
    let %x4 = fn () -> float32 {
      let %x5 = 1f /* ty=float32 */
      %x5
    }
    let %x6 = (%x2 := %x4)
    %x6
  } else {
    let %x7 = fn () -> float32 {
      let %x8 = 2f /* ty=float32 */
      %x8
    }
    let %x9 = (%x2 := %x7)
    %x9
  }
  let %x10 = %x2^
  let %x11 = %x10()
  %x11
}
'''

unit = relay.Function([], relay.const(0., dtype='float32'))
f_ref = relay.Var("f_ref")

cond = relay.var(shape=(), dtype='uint1', name_hint='cond')
true_branch = relay.RefWrite(f_ref, relay.Function([], relay.const(1., dtype='float32')))
false_branch = relay.RefWrite(f_ref, relay.Function([], relay.const(2., dtype='float32')))
if_expr = relay.If(cond, true_branch, false_branch)

stmt = relay.Let(f_ref, relay.RefCreate(unit),
                 relay.Let(relay.Var("x"), if_expr,
                           relay.Call(relay.RefRead(f_ref), [])))

F = relay.Function([cond], stmt)
print(F)
F = relay.ir_pass.infer_type(F)
print(F)
F = relay.ir_pass.to_a_normal_form(F)
F = relay.ir_pass.infer_type(F)
F = relay.ir_pass.partial_evaluate(F)
print(F)

Ideally, what we want should look like

'''
fn (%cond: bool) -> float32 {
  let %x = fn () -> float32 {
    let %x1 = 0f /* ty=float32 */
    %x1
  }
  let %x3 = if (%cond) {
    let %x4 = fn () -> float32 {
      let %x5 = 1f /* ty=float32 */
      %x5
    }
    %x4
  } else {
    let %x7 = fn () -> float32 {
      let %x8 = 2f /* ty=float32 */
      %x8
    }
    %x7
  }
  let %x10 = %x3
  let %x11 = %x10()
  %x11
}
'''