Skip to main content

Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.

Project description

PyTensor logo

Tests Status Coverage

PyTensor is a fork of Aesara – a Python library that allows one to define, optimize, and efficiently evaluate mathematical expressions involving multi-dimensional arrays.

Features

  • 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

Getting started

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)
#  Add [id A]
#  ├─ ExpandDims{axis=0} [id B]
#  │  └─ True_div [id C]
#  │     ├─ a [id D]
#  │     └─ a [id D]
#  └─ dot [id E]
#     ├─ Add [id F]
#     │  ├─ M [id G]
#     │  └─ ExpandDims{axes=[0, 1]} [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)
# Add [id A] 5
#  ├─ [1.] [id B]
#  └─ CGemv{inplace} [id C] 4
#     ├─ AllocEmpty{dtype='float64'} [id D] 3
#     │  └─ Shape_i{0} [id E] 2
#     │     └─ M [id F]
#     ├─ 1.0 [id G]
#     ├─ Add [id H] 1
#     │  ├─ M [id F]
#     │  └─ ExpandDims{axes=[0, 1]} [id I] 0
#     │     └─ a [id J]
#     ├─ v [id K]
#     └─ 0.0 [id L]

See the PyTensor documentation for in-depth tutorials.

Installation

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

Contributing

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.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

pytensor-2.14.2.tar.gz (3.5 MB view hashes)

Uploaded Source

Built Distributions

pytensor-2.14.2-cp311-cp311-win_amd64.whl (1.4 MB view hashes)

Uploaded CPython 3.11 Windows x86-64

pytensor-2.14.2-cp311-cp311-musllinux_1_1_x86_64.whl (1.9 MB view hashes)

Uploaded CPython 3.11 musllinux: musl 1.1+ x86-64

pytensor-2.14.2-cp311-cp311-musllinux_1_1_i686.whl (1.9 MB view hashes)

Uploaded CPython 3.11 musllinux: musl 1.1+ i686

pytensor-2.14.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.9 MB view hashes)

Uploaded CPython 3.11 manylinux: glibc 2.17+ x86-64

pytensor-2.14.2-cp311-cp311-macosx_10_9_x86_64.whl (1.4 MB view hashes)

Uploaded CPython 3.11 macOS 10.9+ x86-64

pytensor-2.14.2-cp310-cp310-win_amd64.whl (1.4 MB view hashes)

Uploaded CPython 3.10 Windows x86-64

pytensor-2.14.2-cp310-cp310-musllinux_1_1_x86_64.whl (1.9 MB view hashes)

Uploaded CPython 3.10 musllinux: musl 1.1+ x86-64

pytensor-2.14.2-cp310-cp310-musllinux_1_1_i686.whl (1.9 MB view hashes)

Uploaded CPython 3.10 musllinux: musl 1.1+ i686

pytensor-2.14.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.9 MB view hashes)

Uploaded CPython 3.10 manylinux: glibc 2.17+ x86-64

pytensor-2.14.2-cp310-cp310-macosx_10_9_x86_64.whl (1.4 MB view hashes)

Uploaded CPython 3.10 macOS 10.9+ x86-64

pytensor-2.14.2-cp39-cp39-win_amd64.whl (1.3 MB view hashes)

Uploaded CPython 3.9 Windows x86-64

pytensor-2.14.2-cp39-cp39-musllinux_1_1_x86_64.whl (1.7 MB view hashes)

Uploaded CPython 3.9 musllinux: musl 1.1+ x86-64

pytensor-2.14.2-cp39-cp39-musllinux_1_1_i686.whl (1.7 MB view hashes)

Uploaded CPython 3.9 musllinux: musl 1.1+ i686

pytensor-2.14.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.7 MB view hashes)

Uploaded CPython 3.9 manylinux: glibc 2.17+ x86-64

pytensor-2.14.2-cp39-cp39-macosx_10_9_x86_64.whl (1.3 MB view hashes)

Uploaded CPython 3.9 macOS 10.9+ x86-64

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page