print uses in output
diff --git a/torch/csrc/jit/ir.cpp b/torch/csrc/jit/ir.cpp
index a5cff1d..a42a00c 100644
--- a/torch/csrc/jit/ir.cpp
+++ b/torch/csrc/jit/ir.cpp
@@ -23,16 +23,19 @@
return THPUtils_unpackString(name.get());
}
}
-
+std::ostream& operator<<(std::ostream & out, Node & n) {
+ if(auto s = n.cast<Select>())
+ out << "%" << s->base()->unique() << "." << s->offset();
+ else
+ out << "%" << n.unique();
+ return out;
+}
std::ostream& operator<<(std::ostream & out, const node_list & nodes) {
size_t i = 0;
for(auto n : nodes) {
if(i++ > 0)
out << ", ";
- if(auto s = n->cast<Select>())
- out << "%" << s->base()->unique() << "." << s->offset();
- else
- out << "%" << n->unique();
+ out << *n;
}
return out;
}
@@ -49,21 +52,28 @@
std::ostream& operator<<(std::ostream & out, Graph & g) {
out << "graph(" << g.inputs() << ") {\n";
for(auto n : g.nodes()) {
- if(!n->cast<Select>()) //improve readibility by printing selects inline
+ if(!n->cast<Select>()) { //improve readibility by printing selects inline
out << " %" << n->unique() << " = ";
- IR_IF(n,PythonOp)
- out << value->name();
- for (auto& scalar : value->scalar_args) {
- out << " " << scalar;
+ IR_IF(n,PythonOp)
+ out << value->name();
+ for (auto& scalar : value->scalar_args) {
+ out << " " << scalar;
+ }
+ IR_ELSEIF(SimpleMap)
+ out << value->op << "!";
+ IR_ELSE()
+ out << toString(n->kind()) << "??";
+ IR_END()
+ out << "(" << n->inputs() << "), uses = [";
+ size_t i = 0;
+ for(auto u : n->uses()) {
+ if(i++ > 0)
+ out << ", ";
+ out << *u.user << ".i" << u.offset;
}
- out << "(" << value->inputs() << ");\n";
- IR_ELSEIF(Select) // print selects inline to improve readability
- IR_ELSEIF(SimpleMap)
- out << value->op << "!(" << value->inputs() << ");\n";
- IR_ELSE()
- out << toString(n->kind()) << "(" << n->inputs() << ");\n";
- IR_END()
+ out << "];\n";
}
+ }
out << " return (" << g.outputs() << ");\n}\n";
return out;
}
