DOCSEARCHbyAlgolia
Submit Docs

PyTensor Documentation

repository·main·Indexed 20 days ago

https://github.com/pymc-devs/pytensor

An optimizing compiler for evaluating mathematical expressions on CPUs and GPUs, serving as the computational backend for PyMC. PyTensor allows for defining and optimizing expressions involving multi-dimensional arrays and supports compilation to backends including JAX, MLX, Numba, and PyTorch.

Tokens
92.7K
Snippets
247
Records
370
Agent score
70%

What's inside PyTensor

  1. Overview of PyTensor features

    main

    PyTensor is a library for defining, optimizing, and evaluating mathematical expressions involving multi-dimensional arrays. Key capabilities include:

    • NumPy Integration: Seamlessly use numpy.ndarray within PyTensor-compiled functions.
    • Symbolic Differentiation: Efficiently computes derivatives for functions with one or many inputs.
    • Numerical Stability: Includes optimizations for edge cases, such as computing log(1 + x) accurately when x is near zero.
    • High-Performance Execution: Generates dynamic C, JAX, or Numba code to evaluate expressions rapidly.

  2. Overview of the PyTensor compilation pipeline

    main

    When you call pytensor.function, PyTensor transforms a symbolic graph into a callable Python object (using Python, C, Numba, or JAX). The compilation process follows four main stages:

    1. Graph Cloning & Variable Collection: The input graph is cloned, shared variables are discovered, givens substitutions are applied, and updates are wired up.
    2. FunctionGraph Creation & Rewriting: A FunctionGraph is built. The Mode (e.g., DebugMode) provides a FunctionMaker that wraps inputs/outputs, applies graph rewrites (optimizations) via a rewriter, and configures a linker.
    3. Linking to a Virtual Machine (VM): The linker (e.g., CLinker or OpWiseCLinker) converts the FunctionGraph into a VM and input/output containers. This stage handles topological sorting and respects flags like strict (for type checking) and borrow (for output storage recycling).
    4. Function Wrapping: The VM and containers are wrapped in a Function object, providing a standard Python callable interface.

  3. What is `scan` and when to use it

    main

    In PyTensor, scan is the primary mechanism for implementing recurrence and looping. While Python for loops can be used, scan is preferred for several reasons:

    • Symbolic Integration: The number of iterations becomes part of the symbolic graph.
    • Automatic Differentiation: It computes gradients through sequential steps.
    • Performance: It is generally faster than a Python for loop when using a compiled PyTensor function.
    • Memory Efficiency: It can lower overall memory usage by detecting the actual amount of memory needed for the operation.

    scan can represent various operations, including reduction and map (looping over leading dimensions), which are special cases of the general scan form.

  4. Understand PyTensor limitations regarding loops and recursion

    main

    PyTensor has specific constraints on control flow within expression graphs:

    • Loops: while- or for-loops are supported only via the pytensor.scan operation. This imposes restrictions on how the loop body interacts with the rest of the graph.
    • Recursion and Goto: Neither goto statements nor recursion are supported or planned within expression graphs.

  5. Implement the Type contract: filter and values_eq_approx

    main

    When subclassing Type, you must implement the filter method. While Type provides default implementations for other methods, filter is required to define how values are validated or cast to your type.

    Additionally, you may want to implement values_eq_approx to allow for approximate equality testing (e.g., to account for numerical instability or rounding errors during graph rewrites).

    from pytensor.graph.type import Type

class DoubleType(Type):
    def filter(self, x, strict=False, allow_downcast=None):
        # Implementation logic for validating/casting x
        ...

    def values_eq_approx(self, x, y, tolerance=1e-4):
        # Implementation logic for approximate equality
        ...

double = DoubleType()

  6. Understand PyTensor core concepts: Variables, Ops, and Applies

    main

    PyTensor operates on an Expression Graph, which is a directed acyclic graph of nodes representing symbolic functional relationships. The core building blocks are:

    • Variable: The primary data structure you work with. It represents a symbolic value (e.g., x = pt.ivector()).
    • Op (Operation): Defines the type of computation to be performed (e.g., pytensor.tensor.add or indexing x[i]).
    • Apply: Represents the actual application of an Op to one or more input Variables to produce output Variables. It is the realization of a mathematical function applied to symbolic inputs.
    • Type: An attribute of a Variable (.type) that indicates the kinds of values computed for it in a compiled graph.
    import pytensor.tensor as pt

x = pt.ivector()
y = -x**2
# x and y are both Variable instances

  7. How views work in PyTensor Ops

    main

    A "view" is an object that shares memory with its source. Changing the source changes the view. When defining a custom Op, you must notify PyTensor which outputs are views of which inputs using the Op.view_map attribute. This allows PyTensor to manage memory and execution order correctly.

    Important Limitation: Currently, an output can only be a view of a single input. Providing a list of multiple inputs for a single output in view_map is not supported.

    from pytensor.graph.op import Op

myop = Op()
# The first output (index 0) is a view of the first input (index 0)
myop.view_map = {0: [0]}

# The first output is a view of the second input (index 1)
myop.view_map = {0: [1]}

# The second output is a view of the first input
myop.view_map = {1: [0]}

# Multiple outputs can be views of different inputs
myop.view_map = {0: [0], 1: [1]}

# Multiple outputs can be views of the same input
myop.view_map = {0: [0], 1: [0]}

  8. Understand the lifecycle of CType methods in generated code

    main

    The methods defined in a CType are called selectively based on the relationship between Python and C for a specific Variable:

    MethodWhen it is called
    c_initWhen a variable is an output/temporary and needs initialization, but no Python object is provided as input.
    c_extractWhen a variable is an input provided from Python. It converts the py_<name> Python object to a C type.
    c_syncWhen a computed value needs to be communicated back to Python (e.g., the final output of a function).
    c_cleanupWhen the C computation for a variable is finished and the data is no longer needed in C.

    Important Lifecycle Notes:

    • c_sync and c_cleanup may be called in sequence. If you allocate memory in c_init or c_extract, ensure c_sync or c_cleanup handles it correctly (e.g., by setting pointers to NULL to avoid double-freeing).
    • c_cleanup is called immediately after any fail code (from c_extract or an Op) is triggered. Therefore, c_cleanup must not depend on any code or variables that appear after a fail reference in the generated block.
    • Only variables declared in c_declare are visible within the scope of c_cleanup.

  9. Understand the Optimization Database (optdb)

    main

    PyTensor uses an ordered database called optdb to manage graph rewrites. It is an instance of SequenceDB (a subclass of RewriteDatabase) that stores Rewriter or RewriteDatabase objects.

    When compiling functions, PyTensor applies a sequence of rewrites from this database. Rewrites can be filtered and queried using tags (e.g., 'fast_run', 'cxx_only', 'inplace'). This allows PyTensor to enable or disable specific optimizations based on the compilation mode or the available backend (like C++).

  10. Understand PyTensor profile output sections

    main

    A PyTensor profile report is divided into four main sections. Understanding these helps identify whether bottlenecks are in PyTensor's overhead, the compilation process, or specific mathematical operations:

    1. Global Info: Contains the function name (set via the name parameter in pytensor.function), the number of calls, and total time. It also shows time spent in Function.vm.__call__ and thunks (PyTensor overhead), as well as time spent in rewriting (graph optimization) and linking (C code compilation).
    2. Class Info: Aggregated information about the classes involved in the execution.
    3. Op Info: Merges information from Apply nodes that share the same Op. Note that some Ops (like Elemwise) only merge if their parameters are identical.
    4. Apply Node Info: Detailed information about every individual Apply node that ran during execution.

    Optimization Tip: To improve performance, focus on the most time-consuming Ops or Apply nodes. You can optimize them by improving their implementation, providing a C implementation, or using graph rewrites to eliminate them.

  11. Limitations of XTensor: Coordinates

    main

    Currently, pytensor.xtensor does not support xarray coordinates. You cannot perform coordinate-based selection operations like .sel.

    However, XTensor graphs that do not use coordinates are highly compatible with the NumPy-like backend of PyTensor and are optimized during the compilation process.

  12. Explore nested graphs with OpFromGraph nodes

    main

    When using pytensor.compile.builders.OpFromGraph, d3viz treats these operations as nested graphs.

    In the interactive HTML visualization, OpFromGraph nodes appear as distinct entities. You can double-click an OpFromGraph node to expand it and view its internal nested graph, including the correct mapping of its input arguments. You can move the expanded graph by dragging it in the shaded area and close it by double-clicking again. This works recursively for nested OpFromGraph compositions.

    import pytensor
import pytensor.d3viz as d3v

x, y, z = pt.scalars('xyz')
e = pt.sigmoid((x + y + z)**2)
# Define an OpFromGraph
op = pytensor.compile.builders.OpFromGraph([x, y, z], [e])

e2 = op(x, y, z) + op(z, y, x)
f = pytensor.function([x, y, z], e2)

# Visualize the function containing OpFromGraph nodes
d3v.d3viz(f, 'examples/ofg.html')