| 97 | /// set the output node to be its first non-moddim node child |
| 98 | template<typename T> |
| 99 | std::vector<TNode<T> *> getClonedOutputNodes( |
| 100 | common::Node_map_t &node_index_map, |
| 101 | const std::vector<std::shared_ptr<common::Node>> &node_clones, |
| 102 | const std::vector<common::Node_ptr> &output_nodes_) { |
| 103 | std::vector<TNode<T> *> cloned_output_nodes; |
| 104 | cloned_output_nodes.reserve(output_nodes_.size()); |
| 105 | for (auto &n : output_nodes_) { |
| 106 | TNode<T> *ptr; |
| 107 | if (n->getOp() == af_moddims_t) { |
| 108 | // if the output node is a moddims node, then set the output node |
| 109 | // to be the child of the moddims node. This is necessary because |
| 110 | // we remove the moddim node_index_map from the tree later |
| 111 | int child_index = node_index_map[n->m_children[0].get()]; |
| 112 | ptr = static_cast<TNode<T> *>(node_clones[child_index].get()); |
| 113 | while (ptr->getOp() == af_moddims_t) { |
| 114 | ptr = static_cast<TNode<T> *>(ptr->m_children[0].get()); |
| 115 | } |
| 116 | } else { |
| 117 | int node_index = node_index_map[n.get()]; |
| 118 | ptr = static_cast<TNode<T> *>(node_clones[node_index].get()); |
| 119 | } |
| 120 | cloned_output_nodes.push_back(ptr); |
| 121 | } |
| 122 | return cloned_output_nodes; |
| 123 | } |
| 124 | |
| 125 | template<typename T> |
| 126 | void evalMultiple(std::vector<Param<T>> arrays, |