Hi All,
Is there a way to return intermediate values inside an IRModule as an output?
Here is an example. How can I return %1 or %0 as an output? from the IRModule or form relay.Function? @mbrookhart any suggestions?
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)
}
When you construct a relay Function, you can wrap multiple output variables in relay.Tuple and return that as the output of the Function.
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
Try this:
new_body = self.visit(fn.body)
print("Visited all", 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)
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
return this super().visit_call(call)
Ohh, I thought my bad! Thanks a lot. There is only a slight problem.
This is what I get. It adds an additional statement into relay IR ( %3 = add(%x1, %y1)
and returns the result.
fn (%x1: Tensor[(1, 1, 1, 20), float32], %y1: Tensor[(1, 1, 1, 20), float32]) {
%0 = add(%x1, %y1);
%1 = subtract(%x1, %y1);
%2 = multiply(%0, %1);
%3 = add(%x1, %y1);
(%2, %3)
}
This is ideally, what I’d like to create. @mbrookhart – Is there something that I can do to fix this?
fn (%x1: Tensor[(1, 1, 1, 20), float32], %y1: Tensor[(1, 1, 1, 20), float32]) {
%0 = add(%x1, %y1);
%1 = subtract(%x1, %y1);
%2 = multiply(%0, %1);
(%2, %0)
}
post = super().visit_call(call)
if post.op.name =="add":
# save this call for return!
self.return_values.append(post)
return post
You were storing the add pre-mutation.
Awesome! Thanks and I believe it works.