Open slyubomirsky opened 1 year ago
We now seem to have the restriction that the body of function must be LeafExpr, IIRC. Probably we can improve it with a better error msg.
If you mean the branches of If
nodes, that would be a mistake if we did it. We require those to be SeqExpr
s.
Edit: If you mean that return
statements should have a leaf expression, I also don't think that's something we should do, as the normalizer should be work.
The restriction is that return is not allowed in true/false branch of if. Not related to recursion. @slyubomirsky a workaround as below.
import tvm
from tvm.script import relax as R
@tvm.script.ir_module
class ControlFlowExample:
@R.function
def a(x: R.Object) -> R.Object:
y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
if y:
r = b(x)
else:
r = c(x)
return r
@R.function
def b(x: R.Object) -> R.Object:
return a(x)
@R.function
def c(x: R.Object) -> R.Object:
return a(x)
If we would like to support multiple return like the following example in Relax Function, one option is to introduce a tag member like Bool return_body
in SeqExprNode
, the tag return_body
is used to tell if body
of SeqExpr
is a return. Any comments? cc: @tqchen @YuchenJin
@R.function
def foo(x: R.Tensor) -> R.Tensor:
y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
if y:
return R.add(x, x)
else:
return R.multiply(x, x)
To keep IR constrained for now, let us not update the IR. this can be sugared to assigning values to global then return. So the only question is do we want to support this syntax in parser and sugar to that IR
Thanks TQ! I don't have a preference on if we need this syntax, a better err message should be good for me.
Okay, thanks for the update. I think the workaround is fine for my purposes. It would be nice for the parser to have syntactic sugar for it and at minimum a better error message.
@tqchen @slyubomirsky I have updated the error message in https://github.com/apache/tvm/pull/14123. I agree it would be nice to have this syntactic sugar (more pythonic style). I can work on adding this sugar.
To keep IR constrained for now, let us not update the IR. this can be sugared to assigning values to global then return. So the only question is do we want to support this syntax in parser and sugar to that IR
What if the if
part returns, but the else
part doesn't?
What if the
if
part returns, but theelse
part doesn't?
Then the code should look like this:
@R.function
def foo(x: R.Tensor) -> R.Tensor:
y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
if y:
return R.add(x, x)
else:
r = R.multiply(x, x)
return r
Shouldn't we first define what return
actually does? The draft of the Relax spec is silent on this.
Shouldn't we first define what
return
actually does? The draft of the Relax spec is silent on this.
return is not in the IR explicitly, currently by default the body
of SeqExpr of function works as return
Yeah, return
is not part of Relax's grammar; it's from Python's grammar and it is up to us to determine how to parse it.
This is actually opening a pretty tough can of worms for parsing, now that I think of it. We could indeed make our lives simpler by requiring both branches to contain a return if either does.
Yeah,
return
is not part of Relax's grammar; it's from Python's grammar and it is up to us to determine how to parse it.This is actually opening a pretty tough can of worms for parsing, now that I think of it. We could indeed make our lives simpler by requiring both branches to contain a return if either does.
Another case is it has no false branch, it may require the member false_branch
of IfNode
to be optional
@R.function
def foo(x: R.Tensor) -> R.Tensor:
y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
if y:
return R.add(x, x)
# no false branch
return R.multiply(x, x)
The Relax AST does not permit omitting the false branch, so how would we parse it? (Even without a return)
The Relax AST does not permit omitting the false branch, so how would we parse it? (Even without a return)
Omitting the false branch is not allowed now, a SeqExpr will be constructed for the false branch during parsing.
Yeah,
return
is not part of Relax's grammar; it's from Python's grammar and it is up to us to determine how to parse it.
Hmm.. So in the example below, does the return
return from the if node, or from the function?
@R.function
def foo(x: R.Tensor) -> R.Tensor:
y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
if y:
return R.add(x, x)
else:
return R.mul(x, x)
...
I think we should parse that to If(y, R.add(x, x), R.mul(x, x))
. The tricky thing is we're parsing Python statements into Relax expressions.
Maybe the easiest thing to do would be to disallow return
at all, except if it's the last statement in the function body, and it's "top-level", i.e. not a proper sub-statement of any other statement.
Edit: The motivation here is that the behavior of return
in the script should match the behavior of it in python code.
Hmm.. So in the example below, does the
return
return from the if node, or from the function?
from the if node.
Edit: The motivation here is that the behavior of
return
in the script should match the behavior of it in python code.
Yeah, this is a bit of the dilemma we run into with this parser. We're trying to fit a square peg (Python's statement-based grammar) into a round hole (Relax's expression-based grammar). It would be nice for common patterns from Python to work smoothly, but we do have room to say that we don't support some things. I'm fine with asking for return values to be outside the if
/else
block, for what it's worth.
I encounter a parse error on the following test case:
The resulting error complains of a violated invariant in the parser: