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