[Relay][Tensorflow] Allow an op as loop var.

[lixiaoquan]

Allow binding a CallNode to a var to support op as loop var.

@zhiics @srkreddy1238 @jroesch Could you please review?

[srkreddy1238]

Look good to me. @jroesch can you too confirm on the IR part ?

[zhiics]

Thanks for the contribution. But I am not sure if we really want to do this because it looks to me that ANF should be able to avoid the mentioned problem. @jroesch please confirm.

[jroesch]

@zhiics I think this is fine, we shouldn't need to make any changes to bind though, the original code was just generating a single loop variable for each sub-graph used as a loop variable this time. We can later ANF the program too.

[jroesch]

We shouldn't need these changes we don't really want to do arbitrary expression to expression replacement.

[jroesch]

I think instead of calling bind which has specific semantics we should just add code like this:

class RewriteSubgraph(ExprMutator):
    def __init__(self, rewrite_map): 
        self.rewrite_map = rewrite_map

    def visit(self, expr):
       if expr in self.rewrite_map:
            return self.rewrite_map[expr]
       else: 
            return super().visit(expr)

def rewrite_subgraph(expr, rewrites):
     return RewriteSubgraph(rewrites).visit(expr)

then replace the call with bind to this, this will handle the general case and not just calls.

[lixiaoquan]

Thanks. Do you think it should be a common function or just be kept in TF frontend?

[jroesch]

I think it is reasonable to just put it in the TF pass for now. We could also put in the frontend/common file where utilities are kept.

[lixiaoquan]

I've modified it according to request.

[jroesch]

Let me know when its ready for review agian.

[lixiaoquan]

@jroesch Will you please review latest change?

[jroesch]

Great looks good, thanks!