From 78671207db483567021a935d0738ba85d8b16551 Mon Sep 17 00:00:00 2001
From: Francois Fleuret <francois@fleuret.org>
Date: Sun, 15 Jan 2017 12:04:13 +0100
Subject: [PATCH] In DAG:saveDot, groups input and output nodes in nested
 clusters.

---
 dagnn.lua | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/dagnn.lua b/dagnn.lua
index 1f45b2a..ca26926 100755
--- a/dagnn.lua
+++ b/dagnn.lua
@@ -185,10 +185,28 @@ end
 function DAG:saveDot(filename)
    local file = (filename and io.open(filename, 'w')) or io.stdout
 
+   local function writeNestedCluster(prefix, list, indent)
+      local indent = indent or ''
+      if torch.type(list) == 'table' then
+         file:write(indent .. '  subgraph cluster_' .. prefix .. ' {\n');
+         for k, x in pairs(list) do
+            writeNestedCluster(prefix .. '_' .. k, x, '  ' .. indent)
+         end
+         file:write(indent .. '  }\n');
+      else
+         file:write(indent .. '  ' .. self.node[list].index .. ' [color=red]\n')
+      end
+   end
+
    file:write('digraph {\n')
 
    file:write('\n')
 
+   writeNestedCluster('input', self.inputModules)
+   writeNestedCluster('output', self.outputModules)
+
+   file:write('\n')
+
    for nnmb, node in pairs(self.node) do
       file:write(
          '  '
-- 
2.20.1