Thanks, @mbrookhart. I think that is promising. However, now my "new_body"
becomes "None" which causes a (surely) error. This is caused by my version
over-written visit_call() function as follows:
I was hoping to save the "result" of add operator to self.return_values with
the following code and use it in the return. This is now causing the
self.visit(fn.body) to return "None". Any ideas?
This is my attempt to save the result of add to a list in order to use it to
wrap the output of a relay.Function.
```
def visit_call(self, call):
if call.op.name =="add":
# save this call for return!
self.return_values.append(call)
super().visit_call(call)
```
Here is the current complete code for your references.
```
class ReWriteOutput(ExprMutator):
"""This pass partitions the subgraph based on the if conditin
"""
def __init__(self):
super().__init__()
# 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 call for return!
self.return_values.append(call)
super().visit_call(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 %1 which is result of subract
"""
new_body = self.visit(fn.body)
print("Visited all! New body: ", new_body)
return_values_to_function = relay.Tuple([new_body] + self.return_values)
func = relay.Function(fn.params, return_values_to_function,
fn.ret_type, fn.type_params, fn.attrs)
print(func)
return func
```
[quote="mbrookhart, post:4, topic:9307"]
Try this:
```
```
[/quote]
---
[Visit
Topic](https://discuss.tvm.apache.org/t/changing-return-of-relay-function-or-irmodule/9307/5)
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/63b94661f85e5a847f18788210c46191860e6b88e0400fd4f9d9cfd394f11465).
