How to get tensor element in relay?

Hi all,

I meet with some problems when I am building my relay graph through tvm relay api.I want to build a graph with a runtime determined control flow. Which looks like this:

x = relay.var("x", shape=(5,), dtype="float32")
one = relay.const(1, dtype="float32")
two = relay.const(2, dtype="float32")
v1 = relay.add(x, one)

v2 = relay.equal(v1[0], two) #this is not legal

true_branch = relay.multiply(two, two)
false_branch = relay.multiply(one, one)
body = relay.If(v2, true_branch, false_branch)

I want to choose the right path depending on the first element of the tensor v1,however relay.equal only accept 0 dimensions arg.Even if I use relay.take to get the first element of v1,the return type is also tensortype with 1 dimension,so it’s not working.

What should I do to make this dynamic control flow working?

Michael