Hi @mbrookhart
Could you please elaborate how I wrap multiple outputs for relay.Fuction?
This is my attempt to add an additional "return" to a relay function. In the
visit_function(), I tried to reconstruct the function, but the relay.Function
prototype does not have any place that I can wrap the returns?
```
class ReWriteOutput(ExprMutator):
"""This pass partitions the subgraph based on the if conditin
"""
def __init__(self):
super().__init__()
self.inputs = []
# self.outputs = relay.Tuple([relay.const(np.random.rand(1, 1)),
relay.const(np.random.rand(1, 1))])
self.return_values = []
def visit_call(self, call):
if call.op.name =="add":
#save_this_for_return
self.return_values.append(call)
super().visit_op(call)
def visit_function(self, fn):
"""Construct a function and apend additional output to the function
The additional input must be one of the intermediate values calculated
inside the
relay function.
Example:
Given this IRModule (or relay Function).
def @main(%x1: Tensor[(1, 1, 1, 20), float32], %y1: Tensor[(1, 1,
1, 20), float32]) {
%0 = add(%x1, %y1);
%1 = subtract(%x1, %y1);
multiply(%0, %1)
}
Modify it so that it returns additional output which in this case
%1 which is the result of subtracting
"""
return_values_to_function = relay.Tuple(self.return_values)
new_body = self.visit(fn.body)
print("Visited all", new_body)
func = relay.Function(fn.params, new_body, fn.ret_type, fn.type_params,
fn.attrs)
return func
```
[quote="mbrookhart, post:2, topic:9307, full:true"]
When you construct a relay Function, you can wrap multiple output variables in
relay.Tuple and return that as the output of the Function.
[/quote]
---
[Visit
Topic](https://discuss.tvm.apache.org/t/changing-return-of-relay-function-or-irmodule/9307/3)
to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click
here](https://discuss.tvm.apache.org/email/unsubscribe/9c146f114550a61b90e8a18afb312ec3d19a0d4eace7c63cab7020cfbd021776).
