text: #include <torch/csrc/jit/passes/lower_tuples.h>
  #include <ATen/core/functional.h>
  #include <c10/util/Exception.h>
  #include <torch/csrc/jit/ir/constants.h>
  #include <torch/csrc/jit/passes/dead_code_elimination.h>
  
  namespace torch {
  namespace jit {
  
  namespace {
  
  // operators where we expect to find tuples as inputs/outputs
  // this is to assert we are only doing modifications when we know
  // we can flatten tuples
  std::unordered_set<Symbol> supported_ops = {
      prim::If,
      prim::Loop,
      prim::TupleUnpack,
      prim::TupleConstruct,
      prim::TupleIndex,
      prim::TupleSlice,
      prim::Param,
      prim::Return,
      prim::PythonOp,
      aten::format,
  };
  
  void removeTupleNodes(Node* n, bool must_remove_tuples) {
    if (n->kind() != prim::TupleUnpack && n->kind() != prim::TupleIndex &&
        n->kind() != prim::TupleSlice) {
      return;
    }
    // tuple index has two inputs, tuple and index
    auto construct = n->inputs().at(0)->node();
    if (construct->kind() != prim::TupleConstruct) {
      if (must_remove_tuples) {
        AT_ERROR(n->kind().toQualString(), " not matched to tuple construct");
      }
      return;
    }
    if (n->kind() == prim::TupleUnpack) {
      for (size_t i = 0; i < n->outputs().size(); ++i) {
        n->outputs()[i]->replaceAllUsesWith(construct->inputs().at(i));
      }
    } else if (n->kind() == prim::TupleIndex) {
      auto idx = n->inputs().at(1);
      auto maybe_int = constant_as<int64_t>(idx);
      if (!maybe_int) {
        if (must_remove_tuples) {
          AT_ERROR(n->sourceRange(), "tuple index with non-constant index");
        }
        return;
      }
      auto int_idx = *maybe_int;
      auto len = construct->output()->type()->containedTypes().size();
      if (int_idx < 0) {
        int_idx += len;
      }
      // currently, we allow non-constant tuple index if the tuple is of one type.
      // so we need to check bounds here
      if (int_idx >= 0 && static_cast<size_t>(int_idx) < len) {
        n->output()->replaceAllUsesWith(construct->inputs().at(int_idx));
      }
    } else if (n->kind() == prim::TupleSlice) {
      std::vector<Value*> values;
      int64_t beg = n->i(attr::beg);
      int64_t end = n->i(attr::end);
      for (int64_t i = beg; i < end; i += 1) {
        values.push_back(construct->inputs().at(i));
      }
      auto graph = n->owningGraph();
      auto tuple_out = graph->createTuple(values);
      WithInsertPoint insert(n);
      graph->insertNode(tuple_out);
      n->output()->replaceAllUsesWith(tuple_out->output());
    }
  }
  } // anonymous namespace
  
  static void LowerAllTuples(Block* block);
  
  static void RemoveTupleConstants(Node* n) {
    if (!(n->kind() == prim::Constant &&
          n->output()->type()->cast<TupleType>())) {
      return;
    }
  
    auto g = n->owningGraph();
    auto tuple_elements = toIValue(n->output()).value().toTuple()->elements();
    WithInsertPoint insert(n);
    std::vector<Value*> elements;
    for (const auto& elem : tuple_elements) {
      auto constant = insertConstant(*n->owningGraph(), elem);
      elements.push_back(constant);
    }
    auto tuple_type = n->output()->type()->expect<TupleType>();
    auto tuple_construct = g->insertNode(n->owningGraph()->createTuple(
        elements, tuple_type->schema() ? tuple_type : nullptr));
  
    // insert the tuple first before recursing on its elements, so that its
    // elements will have a use
    for (Value* elem : elements) {
      RemoveTupleConstants(elem->node());
    }
  
    n->replaceAllUsesWith(tuple_construct);
  }
  
  static void VisitNode(Node* n, Node* insert_point) {
    auto& graph = *n->owningGraph();
  
    // tuple construction operators will become dead when the unpacks are replaced
    if (n->kind() == prim::TupleConstruct) {
      return;
    }
  
    // note: changing the second argument to false changes this pass from a
    // complete lowering pass to one that removes tuples when possible. When
    // tuples are first-class in the interpreter, we should still run this pass to
    // remove extraneous uses
  
    if (n->kind() == prim::TupleUnpack || n->kind() == prim::TupleIndex ||
        n->kind() == prim::TupleSlice) {
      removeTupleNodes(n, /*must_remove_tuples*/ true);
      return;
    }
  
    // flatten the input list  op(a, tup, b) --> op(a, t0, t1, b)
    for (size_t i = 0; i < n->inputs().size();) {
      auto input = n->inputs()[i];
      if (TupleTypePtr tt = input->type()->cast<TupleType>()) {
        TORCH_CHECK(
            supported_ops.count(n->kind()) > 0,
            "tuple appears in op that does not forward tuples, ",
            "unsupported kind: ",
            n->kind().toQualString());
        TORCH_CHECK(
            input->node()->kind() == prim::TupleConstruct,
            "tuple use not matched to tuple construct");
        for (size_t j = 0; j < tt->elements().size(); ++j) {
          n->insertInput(i + 1 + j, input->node()->inputs().at(j));
        }
        n->removeInput(i);
        // note: no update to i
        // since tuples might be nested we need to recursively scan
        // the new flattened inputs
      } else {
        ++i;
      }
    }
    for (auto b : n->blocks()) {
      LowerAllTuples(b);
    }
  
    // flatten the outputs list
    for (size_t i = 0; i < n->outputs().size();) {
      Value* output = n->outputs()[i];
      if (!output->hasUses()) {
        return;
      }
  
      // (a, b, tup, c) -> (a, b, t0, t1, c)
      // and:
      //    tup = (t0, t1)
      // is placed at the current insertion point
      if (TupleTypePtr tt = output->type()->cast<TupleType>()) {
        TORCH_CHECK(
            supported_ops.count(n->kind()) > 0,
            "tuple appears in op that does not forward tuples, ",
            "unsupported kind: ",
            n->kind().toQualString());
        for (size_t j = 0; j < tt->elements().size(); j++) {
          n->insertOutput(i + 1 + j)->setType(tt->elements()[j]);
        }
        auto new_tup =
            graph.createTuple(n->outputs().slice(i + 1, tt->elements().size()));
        new_tup->insertBefore(insert_point);
        insert_point = new_tup;
        output->replaceAllUsesWith(new_tup->output());
        n->eraseOutput(i);
        // note: no update to i to handle nested tuples
      } else {
        ++i;
      }
    }
  }
  
  static void LowerAllTuples(Block* block) {
    // tuples in parameter lists of a block behave exactly the same as
    // _outputs_ of normal instructions, since the param_node represents the
    // parameters as outputs, we can handle it by simply visiting the node
    VisitNode(block->param_node(), *block->nodes().begin());
    for (auto it = block->nodes().begin(), end = block->nodes().end();
         it != end;) {
      auto n = *it++;
      RemoveTupleConstants(n);
      VisitNode(n, *it);
    }
    // tuples in return lists of blocks behave exactly the same as
    // _inputs_ of normal instructions, so we can use VisitNode here as well
    // insert_point is null because it will never be used since return nodes
    // have no outputs
    VisitNode(block->return_node(), nullptr);
  }
  
  static void EnsureNoTuples(ArrayRef<Value*> values) {
    for (Value* v : values) {
      TORCH_CHECK(
          v->type()->kind() != TypeKind::TupleType, "Couldn't lower all tuples.");
    }
  }
  
  static void EnsureNoTuples(Block* block) {
    for (Node* n : block->nodes()) {
      for (Block* b : n->blocks()) {
        EnsureNoTuples(b);
      }
      EnsureNoTuples(n->outputs());
    }
  }
  
  void LowerAllTuples(const std::shared_ptr<Graph>& graph) {
    LowerAllTuples(graph->block());
    EliminateDeadCode(graph->block());
    EnsureNoTuples(graph->block());
  }
  
  void LowerSimpleTuples(Block* block) {
    for (auto n : block->nodes()) {
      removeTupleNodes(n, /*must_remove_tuples*/ false);
      for (auto b : n->blocks()) {
        LowerSimpleTuples(b);
      }
    }
  }
  
  void LowerSimpleTuples(const std::shared_ptr<Graph>& graph) {
    LowerSimpleTuples(graph->block());
    EliminateDeadCode(graph);
  }
  } // namespace jit
  } // namespace torch
  