op – Objects that define operations#

class pytensor.graph.op.HasInnerGraph[source]#

A mixin for an Op that contain an inner graph.

abstract clone()[source]#

Clone the Op and its inner-graph.

fgraph[source]#

The inner function graph (FunctionGraph or FrozenFunctionGraph).

abstract property fn[source]#

The compiled inner-graph function.

abstract property inner_inputs[source]#

The inner function’s inputs.

abstract property inner_outputs[source]#

The inner function’s outputs.

class pytensor.graph.op.Op[source]#

A class that models and constructs operations in a graph.

A Op instance has several responsibilities:

  • construct Apply nodes via Op.make_node() method,

  • perform the numeric calculation of the modeled operation via the Op.perform() method,

  • and (optionally) build the gradient-calculating sub-graphs via the Op.grad() method.

To see how Op, Type, Variable, and Apply fit together see the page on graph – Interface for the PyTensor graph.

For more details regarding how these methods should behave: see the Op Contract in the sphinx docs (advanced tutorial on Op making).

L_op(inputs, outputs, output_grads)[source]#

Construct a graph for the L-operator.

Deprecated since version Implement: pullback() instead.

The L-operator computes a row vector times the Jacobian.

This method dispatches to pullback() if overridden by a subclass, otherwise falls back to Op.grad().

Parameters:
  • inputs – The inputs of the Apply node using this Op.

  • outputs – The outputs of the Apply node using this Op

  • output_grads – The gradients with respect to each Variable in inputs.

R_op(inputs, eval_points)[source]#

Construct a graph for the R-operator.

Deprecated since version Implement: pushforward() instead.

This method is primarily used by Rop. It dispatches to pushforward() if overridden by a subclass.

Parameters:
  • inputs – The Op inputs.

  • eval_points – A Variable or list of Variables with the same length as inputs. Each element of eval_points specifies the value of the corresponding input at the point where the R-operator is to be evaluated.

Return type:

rval[i] should be Rop(f=f_i(inputs), wrt=inputs, eval_points=eval_points).

static add_tag_trace(thing, user_line=None)[source]#

Add tag.trace to a node or variable.

The argument is returned after being affected (inplace).

Parameters:
  • thing – The object where we add .tag.trace.

  • user_line – The max number of user line to keep.

Notes

We also use config.traceback__limit for the maximum number of stack level we look.

default_output = None[source]#

An int that specifies which output Op.__call__() should return. If None, then all outputs are returned.

A subclass should not change this class variable, but instead override it with a subclass variable or an instance variable.

destroy_map = {}[source]#

A dict that maps output indices to the input indices upon which they operate in-place.

Examples

destroy_map = {0: [1]} # first output operates in-place on second input
destroy_map = {1: [0]} # second output operates in-place on first input
do_constant_folding(fgraph, node)[source]#

Determine whether constant folding should be performed for the given node.

This allows each Op to determine if it wants to be constant folded when all its inputs are constant. This allows it to choose where it puts its memory/speed trade-off. Also, it could make things faster as constants can’t be used for in-place operations (see *IncSubtensor).

Parameters:
  • fgraph (FunctionGraph) – Function graph to which node belongs. This is passed in case the Op needs to inspect the graph to make its decision.

  • node (Apply) – The node for which the constant folding determination is made.

Returns:

res

Return type:

bool

grad(inputs, output_grads)[source]#

Construct a graph for the gradient with respect to each input variable.

Deprecated since version Implement: pullback() instead.

Each returned Variable represents the gradient with respect to that input computed based on the symbolic gradients with respect to each output. If the output is not differentiable with respect to an input, then this method should return an instance of type NullType for that input.

Using the reverse-mode AD characterization given in [1], for a \(C = f(A, B)\) representing the function implemented by the Op and its two arguments \(A\) and \(B\), given by the Variables in inputs, the values returned by Op.grad represent the quantities \(\bar{A} \equiv \frac{\partial S_O}{A}\) and \(\bar{B}\), for some scalar output term \(S_O\) of \(C\) in

\[\operatorname{Tr}\left(\bar{C}^\top dC\right) = \operatorname{Tr}\left(\bar{A}^\top dA\right) + \operatorname{Tr}\left(\bar{B}^\top dB\right)\]
Parameters:
  • inputs – The input variables.

  • output_grads – The gradients of the output variables.

Returns:

The gradients with respect to each Variable in inputs.

Return type:

grads

References

inplace_on_inputs(allowed_inplace_inputs)[source]#

Try to return a version of self that tries to inplace in as many as allowed_inplace_inputs.

make_node(*inputs)[source]#

Construct an Apply node that represent the application of this operation to the given inputs.

This must be implemented by sub-classes.

Returns:

node – The constructed Apply node.

Return type:

Apply

make_py_thunk(node, storage_map, compute_map, no_recycling, debug=False)[source]#

Make a Python thunk.

Like Op.make_thunk() but only makes Python thunks.

make_thunk(node, storage_map, compute_map, no_recycling, impl=None)[source]#

Create a thunk.

This function must return a thunk, that is a zero-arguments function that encapsulates the computation to be performed by this op on the arguments of the node.

Parameters:
  • node – Something previously returned by Op.make_node().

  • storage_map – A dict mapping Variables to single-element lists where a computed value for each Variable may be found.

  • compute_map – A dict mapping Variables to single-element lists where a boolean value can be found. The boolean indicates whether the Variable’s storage_map container contains a valid value (i.e. True) or whether it has not been computed yet (i.e. False).

  • no_recycling – List of Variables for which it is forbidden to reuse memory allocated by a previous call.

  • impl (str) – Description for the type of node created (e.g. "c", "py", etc.)

Notes

If the thunk consults the storage_map on every call, it is safe for it to ignore the no_recycling argument, because elements of the no_recycling list will have a value of None in the storage_map. If the thunk can potentially cache return values (like CLinker does), then it must not do so for variables in the no_recycling list.

Op.prepare_node() is always called. If it tries 'c' and it fails, then it tries 'py', and Op.prepare_node() will be called twice.

abstract perform(node, inputs, output_storage)[source]#

Calculate the function on the inputs and put the variables in the output storage.

Parameters:
  • node – The symbolic Apply node that represents this computation.

  • inputs – Immutable sequence of non-symbolic/numeric inputs. These are the values of each Variable in node.inputs.

  • output_storage – List of mutable single-element lists (do not change the length of these lists). Each sub-list corresponds to value of each Variable in node.outputs. The primary purpose of this method is to set the values of these sub-lists.

Notes

The output_storage list might contain data. If an element of output_storage is not None, it has to be of the right type, for instance, for a TensorVariable, it has to be a NumPy ndarray with the right number of dimensions and the correct dtype. Its shape and stride pattern can be arbitrary. It is not guaranteed that such pre-set values were produced by a previous call to this Op.perform(); they could’ve been allocated by another Op’s perform method. An Op is free to reuse output_storage as it sees fit, or to discard it and allocate new memory.

prepare_node(node, storage_map, compute_map, impl)[source]#

Make any special modifications that the Op needs before doing Op.make_thunk().

This can modify the node inplace and should return nothing.

It can be called multiple time with different impl values.

Warning

It is the Op’s responsibility to not re-prepare the node when it isn’t good to do so.

pullback(inputs, outputs, cotangents)[source]#

Construct a graph for the vector-Jacobian product (pullback).

Given a function \(f\) implemented by this Op with inputs \(x\) and outputs \(y = f(x)\), the pullback computes \(\bar{x} = \bar{y}^T J\) where \(J\) is the Jacobian \(\frac{\partial f}{\partial x}\) and \(\bar{y}\) are the cotangent vectors (upstream gradients).

This is the core method for reverse-mode automatic differentiation.

If the output is not differentiable with respect to an input, return a variable of type DisconnectedType for that input. If the gradient is not implemented for some input, return a variable of type NullType (see pytensor.gradient.grad_not_implemented() and pytensor.gradient.grad_undefined()).

Parameters:
  • inputs (Sequence[Variable]) – The input variables of the Apply node using this Op.

  • outputs (Sequence[Variable]) – The output variables of the Apply node using this Op.

  • cotangents (Sequence[Variable]) – The cotangent vectors (gradients w.r.t. each output).

Returns:

input_cotangents – The cotangent vectors w.r.t. each input. One Variable per input.

Return type:

list of Variable

pushforward(inputs, outputs, tangents)[source]#

Construct a graph for the Jacobian-vector product (pushforward).

Given a function \(f\) implemented by this Op with inputs \(x\) and outputs \(y = f(x)\), the pushforward computes \(\dot{y} = J \dot{x}\) where \(J\) is the Jacobian \(\frac{\partial f}{\partial x}\) and \(\dot{x}\) are the tangent vectors.

This is the core method for forward-mode automatic differentiation.

If an output is not differentiable with respect to any input, return a variable of type DisconnectedType for that output. Unlike the legacy R_op method, pushforward must never use None to indicate disconnected outputs.

Parameters:
  • inputs (Sequence[Variable]) – The input variables of the Apply node using this Op.

  • outputs (Sequence[Variable]) – The output variables of the Apply node using this Op.

  • tangents (Sequence[Variable]) – The tangent vectors. One per input. A variable of DisconnectedType indicates that the corresponding input is not being differentiated.

Returns:

output_tangents – The tangent vectors w.r.t. each output. One Variable per output.

Return type:

list of Variable

view_map = {}[source]#

A dict that maps output indices to the input indices of which they are a view.

Examples

view_map = {0: [1]} # first output is a view of second input
view_map = {1: [0]} # second output is a view of first input
class pytensor.graph.op.ThunkType(*args, **kwargs)[source]#
pytensor.graph.op.io_connection_pattern(inputs, outputs)[source]#

Return the connection pattern of a subgraph defined by given inputs and outputs.