PyTensor is a fork of Aesara -- a Python library that allows one to define, optimize, and efficiently evaluate mathematical expressions involving multi-dimensional arrays.
- A hackable, pure-Python codebase
- Extensible graph framework suitable for rapid development of custom operators and symbolic optimizations
- Implements an extensible graph transpilation framework that currently provides compilation via C, JAX, and Numba
- Based on one of the most widely-used Python tensor libraries: Theano
import pytensor from pytensor import tensor as pt # Declare two symbolic floating-point scalars a = pt.dscalar("a") b = pt.dscalar("b") # Create a simple example expression c = a + b # Convert the expression into a callable object that takes `(a, b)` # values as input and computes the value of `c`. f_c = pytensor.function([a, b], c) assert f_c(1.5, 2.5) == 4.0 # Compute the gradient of the example expression with respect to `a` dc = pytensor.grad(c, a) f_dc = pytensor.function([a, b], dc) assert f_dc(1.5, 2.5) == 1.0 # Compiling functions with `pytensor.function` also optimizes # expression graphs by removing unnecessary operations and # replacing computations with more efficient ones. v = pt.vector("v") M = pt.matrix("M") d = a/a + (M + a).dot(v) pytensor.dprint(d) # Elemwise{add,no_inplace} [id A] '' # |InplaceDimShuffle{x} [id B] '' # | |Elemwise{true_div,no_inplace} [id C] '' # | |a [id D] # | |a [id D] # |dot [id E] '' # |Elemwise{add,no_inplace} [id F] '' # | |M [id G] # | |InplaceDimShuffle{x,x} [id H] '' # | |a [id D] # |v [id I] f_d = pytensor.function([a, v, M], d) # `a/a` -> `1` and the dot product is replaced with a BLAS function # (i.e. CGemv) pytensor.dprint(f_d) # Elemwise{Add}[(0, 1)] [id A] '' 5 # |TensorConstant{(1,) of 1.0} [id B] # |CGemv{inplace} [id C] '' 4 # |AllocEmpty{dtype='float64'} [id D] '' 3 # | |Shape_i{0} [id E] '' 2 # | |M [id F] # |TensorConstant{1.0} [id G] # |Elemwise{add,no_inplace} [id H] '' 1 # | |M [id F] # | |InplaceDimShuffle{x,x} [id I] '' 0 # | |a [id J] # |v [id K] # |TensorConstant{0.0} [id L]See the PyTensor documentation for in-depth tutorials.
The latest release of PyTensor can be installed from PyPI using pip:
pip install pytensor
Or via conda-forge:
conda install -c conda-forge pytensor
The current development branch of PyTensor can be installed from GitHub, also using pip:
pip install git+https://github.com/pymc-devs/pytensor
We welcome bug reports and fixes and improvements to the documentation.
For more information on contributing, please see the contributing guide.
A good place to start contributing is by looking through the issues here.