Recent Developments in Pytensor, the Successor Package to Theano
We present the latest developments in Pytensor, the successor package to Theano. Pytensor is a package for defining, manipulating, optimizing, and compiling static computational graphs. We especially focus on full graph-to-graph transformations relevant to the goals of a Bayesian/ML workflow. These allow the user to define a single computational graph, which can then be reused in multiple contexts. In the Bayesian workflow, we are able to extract exact expressions for probabilistic inference from a generative sampling model, or automatically marginalize discrete random variables. In a deep-learning workflow, we can automatically remove dropout and normalization layers when compiling a prediction function from a training graph, or replace expensive operations, such as transformers, with specialized forms at compile time. Finally, we show how the same machinery leads naturally to transpilation into compiled languages, via packages like Numba, Jax, and Pytorch