fix compute inline

Author: leeexyz

include/tvm/te/operation.h

@@ -511,6 +511,7 @@ class HybridOpNode : public OperationNode {
     v->Visit("tag", &tag);
     v->Visit("attrs", &attrs);
     v->Visit("inputs", &inputs);
+    v->Visit("outputs", &outputs);
     v->Visit("symbolic_outputs", &symbolic_outputs);
     v->Visit("axis", &axis);
     v->Visit("body", &body);

python/tvm/te/hybrid/parser.py

@@ -312,7 +312,7 @@ def visit_Assign(self, node):
                     "You should bind a pure name to the tensors",
                 )
                 self.add_symbol(node.targets[i].id, Symbol.GlobalBuffer, rhs.output(i))
-                rmap[rhs.outputs[i].op] = rhs.output(i)
+                rmap[rhs.symbolic_outputs[i].op] = rhs.output(i)
             return utils.replace_io(rhs.body, rmap)
 
         _internal_assert(len(node.targets) == 1, "So far only one-valued assignment is supported!")

src/te/operation/extern_op.cc

@@ -90,6 +90,7 @@ Operation ExternOpNode::ReplaceInputs(const Operation& self,
   ICHECK_EQ(self.operator->(), this);
   auto n = make_object<ExternOpNode>(*this);
   n->body = ReplaceTensor(this->body, rmap);
+  n->outputs = Array<Tensor>();
   for (size_t i = 0; i < n->inputs.size(); ++i) {
     Tensor t = n->inputs[i];
     if (rmap.count(t)) {

src/te/operation/hybrid_op.cc

@@ -104,6 +104,7 @@ Operation HybridOpNode::ReplaceInputs(const Operation& self,
   ICHECK_EQ(self.operator->(), this);
   auto n = make_object<HybridOpNode>(*this);
   n->body = te::ReplaceTensor(this->body, rmap);
+  n->outputs = Array<Tensor>();
   for (size_t i = 0; i < n->inputs.size(); ++i) {
     Tensor t = n->inputs[i];
     if (rmap.count(t)) {

src/te/operation/scan_op.cc

@@ -144,6 +144,7 @@ Operation ScanOpNode::ReplaceInputs(const Operation& self,
                                     const std::unordered_map<Tensor, Tensor>& rmap) const {
   ICHECK_EQ(self.operator->(), this);
   auto n = make_object<ScanOpNode>(*this);
+  n->outputs = Array<Tensor>();
   for (size_t i = 0; i < n->init.size(); ++i) {
     if (rmap.count(n->init[i])) {
       n->init.Set(i, rmap.at(n->init[i]));

src/te/operation/tensor_compute_op.cc

@@ -85,6 +85,7 @@ Operation TensorComputeOpNode::ReplaceInputs(const Operation& self,
                                              const std::unordered_map<Tensor, Tensor>& rmap) const {
   ICHECK_EQ(self.operator->(), this);
   auto n = make_object<TensorComputeOpNode>(*this);
+  n->outputs = Array<Tensor>();
   auto intrin = make_object<TensorIntrinNode>(*(this->intrin.operator->()));
   intrin->body = ReplaceTensor(this->intrin->body, rmap);
   if (intrin->reduce_init.defined()) {

src/te/tensor.cc

@@ -79,6 +79,7 @@ String TensorNode::GetNameHint() const {
 }
 
 Tensor Operation::output(size_t n) const {
+  // cache the output tensors if empty
   if ((*this)->outputs.empty()) {
     auto* ptr = static_cast<OperationNode*>(get_mutable());
     size_t num = static_cast<size_t>((*this)->num_outputs());