File tree 2 files changed +4
-3
lines changed
2 files changed +4
-3
lines changed Original file line number Diff line number Diff line change @@ -121,7 +121,8 @@ export default class ComputationalGraph {
121
121
let s = 'digraph g {\n'
122
122
for ( let i = 0 ; i < this . _nodes . length ; i ++ ) {
123
123
const node = this . nodes [ i ]
124
- s += ` l${ i } [label="${ node . layer . constructor . name } \\n${ node . name } "];\n`
124
+ const label = node . layer . constructor . name + ( node . name ? `\\n${ node . name } ` : '' )
125
+ s += ` l${ i } [label="${ label } "];\n`
125
126
for ( const parent of node . parents ) {
126
127
s += ` l${ parent . index } -> l${ i } ;\n`
127
128
}
Original file line number Diff line number Diff line change @@ -158,9 +158,9 @@ describe('Computational Graph', () => {
158
158
test ( 'toDot' , ( ) => {
159
159
const graph = new ComputationalGraph ( )
160
160
graph . add ( Layer . fromObject ( { type : 'input' } ) )
161
- graph . add ( Layer . fromObject ( { type : 'tanh' } ) )
161
+ graph . add ( Layer . fromObject ( { type : 'tanh' } ) , 't' )
162
162
expect ( graph . toDot ( ) ) . toBe (
163
- 'digraph g {\n l0 [label="InputLayer\\nundefined "];\n l1 [label="TanhLayer\\nundefined "];\n l0 -> l1;\n}'
163
+ 'digraph g {\n l0 [label="InputLayer"];\n l1 [label="TanhLayer\\nt "];\n l0 -> l1;\n}'
164
164
)
165
165
} )
166
166
You can’t perform that action at this time.
0 commit comments