Make it possible to override the node color in dot graphs via DotOptions.

Change: 123212554

1 files changed, 5 insertions, 0 deletions

diff --git a/tensorflow/core/graph/dot.h b/tensorflow/core/graph/dot.h
index 79a538978a..96e48773a9 100644
--- a/tensorflow/core/graph/dot.h
+++ b/tensorflow/core/graph/dot.h
@@ -48,6 +48,11 @@ struct DotOptions {
   // A function that returns the "cost" of the edge.  The dot display
   // makes a edge thickness proportional to its cost.
   std::function<double(const Edge*)> edge_cost;
+
+  // A function that returns a color number to apply to each node. < 0 means
+  // no color. A color will be assigned to each color number from a palette;
+  // adjacent color numbers will receive different colors.
+  std::function<int(const Node*)> node_color;
 };
 
 // Return a string that contains a graphviz specification of the graph.