.. _tutorial_printing_drawing:

================================
Printing/Drawing PyTensor graphs
================================


PyTensor provides the functions :func:`pytensor.printing.pprint` and
:func:`pytensor.printing.debugprint` to print a graph to the terminal before or
after compilation. :func:`pprint` is more compact and math-like,
:func:`debugprint` is more verbose. PyTensor also provides :func:`pydotprint`
that creates an image of the function. You can read about them in
:ref:`libdoc_printing`.

.. note::

    When printing PyTensor functions, they can sometimes be hard to
    read.  To help with this, you can disable some PyTensor rewrites
    by using the PyTensor flag:
    ``optimizer_excluding=fusion:inplace``. Do not use this during
    real job execution, as this will make the graph slower and use more
    memory.

Consider again the logistic regression example:

>>> import numpy as np
>>> import pytensor
>>> import pytensor.tensor as pt
>>> rng = np.random.default_rng(2382)
>>> # Training data
>>> N = 400
>>> feats = 784
>>> D = (rng.standard_normal(N, feats).astype(pytensor.config.floatX), rng.integers(size=N,low=0, high=2).astype(pytensor.config.floatX))
>>> training_steps = 10000
>>> # Declare PyTensor symbolic variables
>>> x = pt.matrix("x")
>>> y = pt.vector("y")
>>> w = pytensor.shared(rng.standard_normal(feats).astype(pytensor.config.floatX), name="w")
>>> b = pytensor.shared(np.asarray(0., dtype=pytensor.config.floatX), name="b")
>>> x.tag.test_value = D[0]
>>> y.tag.test_value = D[1]
>>> # Construct PyTensor expression graph
>>> p_1 = 1 / (1 + pt.exp(-pt.dot(x, w)-b)) # Probability of having a one
>>> prediction = p_1 > 0.5 # The prediction that is done: 0 or 1
>>> # Compute gradients
>>> xent = -y*pt.log(p_1) - (1-y)*pt.log(1-p_1) # Cross-entropy
>>> cost = xent.mean() + 0.01*(w**2).sum() # The cost to optimize
>>> gw,gb = pt.grad(cost, [w,b])
>>> # Training and prediction function
>>> train = pytensor.function(inputs=[x,y], outputs=[prediction, xent], updates=[[w, w-0.01*gw], [b, b-0.01*gb]], name = "train")
>>> predict = pytensor.function(inputs=[x], outputs=prediction, name = "predict")


Pretty Printing
===============

>>> pytensor.printing.pprint(prediction) # doctest: +NORMALIZE_WHITESPACE
'gt((TensorConstant{1} / (TensorConstant{1} + exp(((-(x \\dot w)) - b)))),
TensorConstant{0.5})'


Debug Print
===========

The pre-compilation graph:

>>> pytensor.printing.debugprint(prediction) # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
Elemwise{gt,no_inplace} [id A] ''
 |Elemwise{true_div,no_inplace} [id B] ''
 | |InplaceDimShuffle{x} [id C] ''
 | | |TensorConstant{1} [id D]
 | |Elemwise{add,no_inplace} [id E] ''
 |   |InplaceDimShuffle{x} [id F] ''
 |   | |TensorConstant{1} [id D]
 |   |Elemwise{exp,no_inplace} [id G] ''
 |     |Elemwise{sub,no_inplace} [id H] ''
 |       |Elemwise{neg,no_inplace} [id I] ''
 |       | |dot [id J] ''
 |       |   |x [id K]
 |       |   |w [id L]
 |       |InplaceDimShuffle{x} [id M] ''
 |         |b [id N]
 |InplaceDimShuffle{x} [id O] ''
   |TensorConstant{0.5} [id P]

The post-compilation graph:

>>> pytensor.printing.debugprint(predict)  # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
Elemwise{Composite{GT(scalar_sigmoid((-((-i0) - i1))), i2)}} [id A] ''   4
 |...Gemv{inplace} [id B] ''   3
 | |AllocEmpty{dtype='float64'} [id C] ''   2
 | | |Shape_i{0} [id D] ''   1
 | |   |x [id E]
 | |TensorConstant{1.0} [id F]
 | |x [id E]
 | |w [id G]
 | |TensorConstant{0.0} [id H]
 |InplaceDimShuffle{x} [id I] ''   0
 | |b [id J]
 |TensorConstant{(1,) of 0.5} [id K]


Picture Printing of Graphs
==========================

The pre-compilation graph:

>>> pytensor.printing.pydotprint(prediction, outfile="pics/logreg_pydotprint_prediction.png", var_with_name_simple=True)  # doctest: +SKIP
The output file is available at pics/logreg_pydotprint_prediction.png

.. image:: ./pics/logreg_pydotprint_prediction.png
   :width: 800 px

The post-compilation graph:

>>> pytensor.printing.pydotprint(predict, outfile="pics/logreg_pydotprint_predict.png", var_with_name_simple=True)  # doctest: +SKIP
The output file is available at pics/logreg_pydotprint_predict.png

.. image:: ./pics/logreg_pydotprint_predict.png
   :width: 800 px

The optimized training graph:

>>> pytensor.printing.pydotprint(train, outfile="pics/logreg_pydotprint_train.png", var_with_name_simple=True)  # doctest: +SKIP
The output file is available at pics/logreg_pydotprint_train.png

.. image:: ./pics/logreg_pydotprint_train.png
   :width: 1500 px


Interactive Graph Visualization
===============================

The new :mod:`d3viz` module complements :func:`pytensor.printing.pydotprint` to
visualize complex graph structures. Instead of creating a static image, it
generates an HTML file, which allows to dynamically inspect graph structures in
a web browser. Features include zooming, drag-and-drop, editing node labels, or
coloring nodes by their compute time.

=> :mod:`d3viz` <=

.. image:: ./pics/d3viz.png
   :width: 350 px