public
Last active

IPython extension for automatically PNG-pretty-printing Theano functions with Graphviz.

  • Download Gist
theanoprinting.py
Python
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35
"""Display Theano functions in the IPython notebook with pydotprint."""
 
__author__ = "David Warde-Farley"
__copyright__ = "Copyright 2012, Universite de Montreal"
__credits__ = ["David Warde-Farley"]
__license__ = "3-clause BSD"
__email__ = "wardefar@iro"
__maintainer__ = "David Warde-Farley"
 
 
import os
import tempfile
from theano.compile.function_module import Function
from theano.printing import pydotprint
 
_loaded = False
 
 
def print_function_png(o):
handle, fn = tempfile.mkstemp(suffix='.png')
try:
os.close(handle)
pydotprint(o, outfile=fn, format='png', print_output_file=False)
with open(fn) as f:
return f.read()
finally:
os.remove(fn)
 
 
def load_ipython_extension(ip):
global _loaded
if not _loaded:
png_formatter = ip.display_formatter.formatters['image/png']
png_formatter.for_type(Function, print_function_png)
_loaded = True

By the way, we're trying to move away from the global _loaded pattern. New versions of IPython will track whether the extension is already loaded, so the extension itself doesn't need to. In current versions, setting the formatter more than once will make no difference, so there's no need to guard against it happening several times.

The aim of this is to reduce global state, so that we can one day have more than one IPython instance in a process.

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.