`OpFromGraph`#

This page describes pytensor.compile.builders.OpFromGraph, an Op constructor that allows one to encapsulate an PyTensor graph in a single Op.

This can be used to encapsulate some functionality in one block. It is useful to scale PyTensor compilation for regular bigger graphs when we reuse that encapsulated functionality with different inputs many times. Due to this encapsulation, it can make PyTensor’s compilation phase faster for graphs with many nodes.

Using this for small graphs is not recommended as it disables rewrites between what is inside the encapsulation and outside of it.

class pytensor.compile.builders.OpFromGraph(inputs, outputs, *, inline=False, pullback=None, pushforward=None, lop_overrides=None, rop_overrides=None, connection_pattern=None, strict=False, name=None, destroy_map=None, **kwargs)[source]#

Create an Op from inputs and outputs lists of variables.

The signature is similar to pytensor.function() and the resulting Op’s perform will do the same operation as pytensor.function(inputs, outputs, **kwargs).

Does not support updates or givens.

Todo

Add support for NullType and DisconnectedType when R_op supports them
Add optimization to removing unused inputs/outputs
Add optimization to work inplace on inputs when not inline

Notes

Shared variables in the inner graph are supported. They are detected automatically and added as implicit inputs.
Unused inputs are supported (needed for gradient overrides).
Nested OpFromGraph is supported.
inline=True causes the Op’s inner graph to be inlined during compilation, which gives better runtime optimization at the cost of compilation time. Currently only works with fast_compile or fast_run mode.
Override callables should be pure functions (no side effects). They are called once at the first call to L_op/R_op and converted to OpFromGraph instances. They are also called once at construction time with dummy inputs to build a frozen representation for equality comparison.
Two OpFromGraph instances with the same inner graph, overrides, shared variables, and settings are considered equal. This allows the MergeOptimizer to deduplicate identical OpFromGraph nodes.

Examples

Basic usage:

from pytensor import function, tensor as pt
from pytensor.compile.builders import OpFromGraph

x, y, z = pt.scalars("xyz")
e = x + y * z
op = OpFromGraph([x, y, z], [e])
# op behaves like a normal pytensor op
e2 = op(x, y, z) + op(z, y, x)
fn = function([x, y, z], [e2])

With a shared variable:

import numpy as np
import pytensor
from pytensor import config, function, tensor as pt
from pytensor.compile.builders import OpFromGraph

x, y, z = pt.scalars("xyz")
s = pytensor.shared(np.random.random((2, 2)).astype(config.floatX))
e = x + y * z + s
op = OpFromGraph([x, y, z], [e])
e2 = op(x, y, z) + op(z, y, x)
fn = function([x, y, z], [e2])

Per-input L_op override:

from pytensor import function, tensor as pt, grad
from pytensor.compile.builders import OpFromGraph

x, y, z = pt.scalars("xyz")
e = x + y * z

def rescale_dy(inps, outputs, out_grads):
    x, y, z = inps
    (g,) = out_grads
    return z * 2


op = OpFromGraph(
    [x, y, z],
    [e],
    pullback=[None, rescale_dy, None],
)
e2 = op(x, y, z)
dx, dy, dz = grad(e2, [x, y, z])
fn = function([x, y, z], [dx, dy, dz])
# the gradient wrt y is now doubled
fn(2.0, 3.0, 4.0)  # [1., 8., 3.]

OpFromGraph#

`OpFromGraph`#