diff --git a/doc/extending/creating_an_op.rst b/doc/extending/creating_an_op.rst
index e42241b92e..8744b3b537 100644
--- a/doc/extending/creating_an_op.rst
+++ b/doc/extending/creating_an_op.rst
@@ -4,37 +4,15 @@
Creating a new :class:`Op`: Python implementation
=================================================
-So suppose you have looked through the library documentation and you don't see
-a function that does what you want.
+You may have looked through the library documentation but don't see a function that does what you want.
-If you can implement something in terms of an existing :ref:`op`, you should do that.
-Odds are your function that uses existing PyTensor expressions is short,
-has no bugs, and potentially profits from rewrites that have already been
-implemented.
+If you can implement something in terms of an existing :class:`Op`, you should do that.
+A PyTensor function that builds upon existing expressions will be better optimized, automatic differentiable, and
+work seamlessly across different backends.
-However, if you cannot implement an :class:`Op` in terms of an existing :class:`Op`, you have to
-write a new one.
-
-As an illustration, this tutorial will demonstrate how a simple Python-based
-:class:`Op` that performs operations on ``np.float64``\s is written.
-
-.. note::
-
- This is an introductory tutorial and as such it does not cover how to make
- an :class:`Op` that returns a view or modifies the values in its inputs. Thus, all
- :class:`Op`\s created with the instructions described here MUST return newly
- allocated memory or reuse the memory provided in the parameter
- ``output_storage`` of the :meth:`Op.perform` method. See
- :ref:`views_and_inplace` for an explanation on how to do this.
-
- If your :class:`Op` returns a view or changes the value of its inputs
- without doing as prescribed in that page, PyTensor will run, but will
- return correct results for some graphs and wrong results for others.
-
- It is recommended that you run your tests in :class:`DebugMode`, since it
- can help verify whether or not your :class:`Op` behaves correctly in this
- regard.
+However, if you cannot implement an :class:`Op` in terms of an existing :class:`Op`, you have to write a new one.
+This page will show how to implement some simple Python-based :class:`Op` that perform operations on numpy arrays.
PyTensor Graphs refresher
-------------------------
@@ -45,12 +23,12 @@ PyTensor Graphs refresher
PyTensor represents symbolic mathematical computations as graphs. Those graphs
are bi-partite graphs (graphs with two types of nodes), they are composed of
interconnected :ref:`apply` and :ref:`variable` nodes.
-:class:`Variable` nodes represent data in the graph, either inputs, outputs or
+:ref:`variable` nodes represent data in the graph, either inputs, outputs or
intermediary values. As such, inputs and outputs of a graph are lists of PyTensor
-:class:`Variable` nodes. :class:`Apply` nodes perform computation on these
-variables to produce new variables. Each :class:`Apply` node has a link to an
+:ref:`variable` nodes. :ref:`apply` nodes perform computation on these
+variables to produce new variables. Each :ref:`apply` node has a link to an
instance of :class:`Op` which describes the computation to perform. This tutorial
-details how to write such an :class:`Op` instance. Please refers to
+details how to write such an :class:`Op` instance. Please refer to
:ref:`graphstructures` for a more detailed explanation about the graph
structure.
@@ -58,338 +36,263 @@ structure.
:class:`Op`'s basic methods
---------------------------
-An :class:`Op` is any Python object which inherits from :class:`Op`.
+An :class:`Op` is any Python object that inherits from :class:`Op`.
This section provides an overview of the basic methods you typically have to
implement to make a new :class:`Op`. It does not provide extensive coverage of all the
possibilities you may encounter or need. For that refer to
-:ref:`op_contract`.
+:ref:`Op contract <op_contract>`.
.. testcode:: python
- import pytensor
+ from typing import Any
+ from pytensor.graph.basic import Apply, Variable
+ from pytensor.graph.fg import FunctionGraph
from pytensor.graph.op import Op
+ from pytensor.graph.type import Type
class MyOp(Op):
# Properties attribute
- __props__ = ()
+ __props__ : tuple[Any, ...] = ()
- #itypes and otypes attributes are
- #compulsory if make_node method is not defined.
- #They're the type of input and output respectively
- itypes = None
- otypes = None
-
- #Compulsory if itypes and otypes are not defined
- def make_node(self, *inputs):
+ # Constructor, usually used only to set Op properties
+ def __init__(self, *args):
pass
- # Python implementation:
- def perform(self, node, inputs_storage, output_storage):
- pass
+ # itypes and otypes attributes are compulsory if make_node method is not defined.
+ # They're the type of input and output respectively
+ itypes: list[Type] | None = None
+ otypes: list[Type] | None = None
- # Other type of implementation
- # C implementation: [see pytensor web site for other functions]
- def c_code(self, node, inputs, outputs, sub):
+ # make_node is compulsory if itypes and otypes are not defined
+ # make_node is more flexible: output types can be determined
+ # based on the input types and Op properties.
+ def make_node(self, *inputs) -> Apply:
pass
- # Other implementations:
- def make_thunk(self, node, storage_map, _, _2, impl=None):
+ # Performs the numerical evaluation of Op in Python. Required.
+ def perform(self, node: Apply, inputs_storage: list[Any], output_storage: list[list[Any]]) -> None:
pass
- # optional:
- check_input = True
-
- def __init__(self, *args):
+ # Defines the symbolic expression for the L-operator based on the input and output variables
+ # and the output gradient variables. Optional.
+ def L_op(self, inputs: list[Variable], outputs: list[Variable], output_grads: list[Variable]) -> list[Variable]:
pass
- def grad(self, inputs, g):
+ # Equivalent to L_op, but with a "technically"-bad name and without outputs provided.
+ # It exists for historical reasons. Optional.
+ def grad(self, inputs: list[Variable], output_grads: list[Variable]) -> list[Variable]:
+ # Same as self.L_op(inputs, self(inputs), output_grads)
pass
- def R_op(self, inputs, eval_points):
+ # Defines that symbolic expression for the R-operator based on the input variables
+ # and eval_point variables. Optional.
+ def R_op(self, inputs: list[Variable], eval_points: list[Variable | None]) -> list[Variable | None]:
pass
- def infer_shape(self, fgraph, node, input_shapes):
+ # Defines the symbolic expression for the output shape based on the input shapes
+ # and, less frequently, the input variables via node.inputs. Optional.
+ def infer_shape(self, fgraph: FunctionGraph, node: Apply, input_shapes: list[tuple[Variable, ...]]) -> list[tuple[Variable]]:
pass
An :class:`Op` has to implement some methods defined in the the interface of
:class:`Op`. More specifically, it is mandatory for an :class:`Op` to define either
-the method :meth:`Op.make_node` or :attr:`Op.itypes`, :attr:`Op.otypes` and one of the
-implementation methods, either :meth:`Op.perform`, :meth:`COp.c_code`
-or :meth:`Op.make_thunk`.
-
- :meth:`Op.make_node` method creates an Apply node representing the application
- of the :class:`Op` on the inputs provided. This method is responsible for three things:
-
- - it first checks that the input :class:`Variable`\s types are compatible
- with the current :class:`Op`. If the :class:`Op` cannot be applied on the provided
- input types, it must raises an exception (such as :class:`TypeError`).
- - it operates on the :class:`Variable`\s found in
- ``*inputs`` in PyTensor's symbolic language to infer the type of
- the symbolic output :class:`Variable`\s. It creates output :class:`Variable`\s of a suitable
- symbolic :class:`Type` to serve as the outputs of this :class:`Op`'s
- application.
- - it creates an :class:`Apply` instance with the input and output :class:`Variable`, and
- return the :class:`Apply` instance.
-
-
-
- :meth:`Op.perform` method defines the Python implementation of an :class:`Op`.
- It takes several arguments:
-
- - ``node`` is a reference to an Apply node which was previously
- obtained via the :meth:`Op.make_node` method. It is typically not
- used in a simple :class:`Op`, but it contains symbolic information that
- could be required by a complex :class:`Op`.
- - ``inputs`` is a list of references to data which can be operated on using
- non-symbolic statements, (i.e., statements in Python, Numpy).
- - ``output_storage`` is a list of storage cells where the output
- is to be stored. There is one storage cell for each output of the :class:`Op`.
- The data put in ``output_storage`` must match the type of the
- symbolic output. It is forbidden to change the length of the list(s)
- contained in ``output_storage``.
- A function Mode may allow ``output_storage`` elements to persist
- between evaluations, or it may reset ``output_storage`` cells to
- hold a value of ``None``. It can also pre-allocate some memory
- for the :class:`Op` to use. This feature can allow ``perform`` to reuse
- memory between calls, for example. If there is something
- preallocated in the ``output_storage``, it will be of the good
- dtype, but can have the wrong shape and have any stride pattern.
-
- :meth:`Op.perform` method must be determined by the inputs. That is to say,
- when applied to identical inputs the method must return the same outputs.
-
- An :class:`Op`\s implementation can be defined in other ways, as well.
- For instance, it is possible to define a C-implementation via :meth:`COp.c_code`.
- Please refers to tutorial :ref:`creating_a_c_op` for a description of
- :meth:`COp.c_code` and other related ``c_**`` methods. Note that an
- :class:`Op` can provide both Python and C implementations.
-
- :meth:`Op.make_thunk` method is another alternative to :meth:`Op.perform`.
- It returns a thunk. A thunk is defined as a zero-arguments
- function which encapsulates the computation to be performed by an
- :class:`Op` on the arguments of its corresponding node. It takes several parameters:
-
- - ``node`` is the :class:`Apply` instance for which a thunk is requested,
- - ``storage_map`` is a ``dict`` of lists which maps variables to a one-element
- lists holding the variable's current value. The one-element list acts as
- pointer to the value and allows sharing that "pointer" with other nodes
- and instances.
- - ``compute_map`` is also a dict of lists.
- It maps variables to one-element lists holding booleans. If
- the value is 0 then the variable has not been computed and the
- value should not be considered valid. If the value is 1 the
- variable has been computed and the value is valid. If the value
- is 2 the variable has been garbage-collected and is no longer
- valid, but shouldn't be required anymore for this call.
- The returned function must ensure that it sets the computed
- variables as computed in the :obj:`compute_map`.
- - ``impl`` allow to select between multiple implementation.
- It should have a default value of ``None``.
-
- :meth:`Op.make_thunk` is useful if you want to generate code and compile
- it yourself.
-
- If :meth:`Op.make_thunk` is defined by an :class:`Op`, it will be used by PyTensor
- to obtain the :class:`Op`'s implementation.
- :meth:`Op.perform` and :meth:`COp.c_code` will be ignored.
-
- If :meth:`Op.make_node` is not defined, the :attr:`Op.itypes` and :attr:`Op.otypes`
- are used by the :class:`Op`'s :meth:`Op.make_node` method to implement the functionality
- of :meth:`Op.make_node` method mentioned above.
+the method :meth:`make_node` or :attr:`itypes`, :attr:`otypes`, and :meth:`perform`.
+
+:meth:`make_node`
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+:meth:`make_node` method creates an :ref:`apply` node representing the application
+of the :class:`Op` on the inputs provided. This method is responsible for three things:
+
+- Checks that the inputs can be converted to :ref:`variable`\s whose types are compatible with the current :class:`Op`.
+ If the :class:`Op` cannot be applied on the provided input types, it must raises an exception (such as :class:`TypeError`).
+- Creates new output :ref:`variable`\s of a suitable symbolic :class:`Type` to serve as the outputs of this :class:`Op`'s application.
+- Returns an :ref:`apply` instance with the input and output :ref:`variable`\s, and itself as the :class:`Op`.
+
+If :meth:`make_node` is not defined, the :attr:`itypes` and :attr:`otypes` are used by the :class:`Op`'s
+:meth:`make_node` method to implement the functionality method mentioned above.
+
+
+:meth:`perform`
+^^^^^^^^^^^^^^^^^^
+
+:meth:`perform` method defines the Python implementation of an :class:`Op`.
+It takes several arguments:
+
+- ``node`` is a reference to an :ref:`apply` node which was previously
+ obtained via the :meth:`make_node` method. It is typically not
+ used in a simple :class:`Op`, but it contains symbolic information that
+ could be required by a complex :class:`Op`.
+- ``inputs`` is a list of references to data which can be operated on using
+ non-symbolic statements, (i.e., statements in Python, Numpy).
+- ``output_storage`` is a list of storage cells where the output
+ is to be stored. There is one storage cell for each output of the :class:`Op`.
+ The data put in ``output_storage`` must match the type of the
+ symbolic output.
+ PyTensor may sometimes allow ``output_storage`` elements to persist
+ between evaluations, or it may reset ``output_storage`` cells to
+ hold a value of ``None``. It can also pre-allocate some memory
+ for the :class:`Op` to use. This feature can allow ``perform`` to reuse
+ memory between calls, for example. If there is something
+ preallocated in the ``output_storage``, it will be of the correct
+ dtype, but can have the wrong shape and have any stride pattern.
+
+:meth:`perform` method must be determined by the inputs.
+That is to say, when applied to identical inputs the method must return the same outputs.
+
:class:`Op`'s auxiliary methods
-------------------------------
There are other methods that can be optionally defined by the :class:`Op`:
- :meth:`Op.__eq__` and :meth:`Op.__hash__` define respectively equality
- between two :class:`Op`\s and the hash of an :class:`Op` instance.
- They will be used during the rewriting phase to merge nodes that are doing
- equivalent computations (same inputs, same operation).
- Two :class:`Op`\s that are equal according :meth:`Op.__eq__`
- should return the same output when they are applied on the same inputs.
-
- The :attr:`Op.__props__` attribute lists the properties that influence how the computation
- is performed. Usually these are set in :meth:`Op.__init__`. It must be a tuple.
- If you don't have any properties, then you should set this attribute to the
- empty tuple ``()``.
-
- :attr:`Op.__props__` enables the automatic generation of appropriate
- :meth:`Op.__eq__` and :meth:`Op.__hash__`.
- Given the method :func:`__eq__`, automatically generated from
- :attr:`Op.__props__`, two :class:`Op`\s will be equal if they have the same values for all
- the properties listed in :attr:`Op.__props__`.
- Given to the method :meth:`Op.__hash__` automatically generated from
- :attr:`Op.__props__`, two :class:`Op`\s will be have the same hash if they have the same
- values for all the properties listed in :attr:`Op.__props__`.
- :attr:`Op.__props__` will also generate a suitable :meth:`Op.__str__` for your :class:`Op`.
-
- The :meth:`Op.infer_shape` method allows an :class:`Op` to infer the shape of its
- output variables without actually computing them.
- It takes as input ``fgraph``, a :class:`FunctionGraph`; ``node``, a reference
- to the :class:`Op`'s :class:`Apply` node;
- and a list of :class:`Variables`\s (e.g. ``i0_shape``, ``i1_shape``, ...)
- which are the dimensions of the :class:`Op` input :class:`Variable`\s.
- :meth:`Op.infer_shape` returns a list where each element is a tuple representing
- the shape of one output.
- This could be helpful if one only needs the shape of the output instead of the
- actual outputs, which can be useful, for instance, for rewriting
- procedures.
-
- The :meth:`Op.grad` method is required if you want to differentiate some cost
- whose expression includes your :class:`Op`. The gradient may be
- specified symbolically in this method. It takes two arguments ``inputs`` and
- ``output_gradients``, which are both lists of :class:`Variable`\s, and
- those must be operated on using PyTensor's symbolic language. The :meth:`Op.grad`
- method must return a list containing one :class:`Variable` for each
- input. Each returned :class:`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 be defined to return a variable of type :class:`NullType`
- for that input. Likewise, if you have not implemented the gradient
- computation for some input, you may return a variable of type
- :class:`NullType` for that input. Please refer to :meth:`Op.grad` for a more detailed
- view.
-
- The :meth:`Op.R_op` method is needed if you want :func:`pytensor.gradient.Rop` to
- work with your :class:`Op`.
- This function implements the application of the R-operator on the
- function represented by your :class:`Op`. Let assume that function is :math:`f`,
- with input :math:`x`, applying the R-operator means computing the
- Jacobian of :math:`f` and right-multiplying it by :math:`v`, the evaluation
- point, namely: :math:`\frac{\partial f}{\partial x} v`.
-
- The optional boolean :attr:`check_input` attribute is used to specify
- if you want the types used in your :class:`COp` to check their inputs in their
- :meth:`COp.c_code`. It can be used to speed up compilation, reduce overhead
- (particularly for scalars) and reduce the number of generated C files.
+:attr:`__props__`
+^^^^^^^^^^^^^^^^^^^^
+The :attr:`__props__` attribute lists the :class:`Op` instance properties
+that influence how the computation is performed. It must be a hashable tuple.
+Usually these are set in :meth:`__init__`. If you don't have any properties
+that influence the computation, then you will want to set this attribute to the empty tuple ``()``.
-Example: :class:`Op` definition
--------------------------------
+:attr:`__props__` enables the automatic generation of appropriate :meth:`__eq__` and :meth:`__hash__`.
+According to this default, :meth:`__eq__`, two :class:`Op`\s will be equal if they have the same values for all
+the properties listed in :attr:`__props__`. Similarly, they will have the same hash.
-.. testcode:: example
+When PyTensor sees two nodes with equal :class:`Op`\s and the same set of inputs,
+it will assume the outputs are equivalent and merge the nodes to avoid redundant computation.
+When `Op.__props__` is not specified, two distinct instances of the same class will not be equal
+and hash to their `id`. PyTensor won't merge nodes with the same class but different instances in this case.
- import pytensor
- from pytensor.graph.op import Op
- from pytensor.graph.basic import Apply
+:attr:`__props__` will also generate a suitable :meth:`__repr__` and :meth:`__str__` for your :class:`Op`.
- class DoubleOp1(Op):
- __props__ = ()
+:meth:`infer_shape`
+^^^^^^^^^^^^^^^^^^^^^^
- def make_node(self, x):
- x = pytensor.tensor.as_tensor_variable(x)
- # Note: using x_.type() is dangerous, as it copies x's broadcasting
- # behaviour
- return Apply(self, [x], [x.type()])
+The :meth:`infer_shape` method allows an :class:`Op` to infer the shape of its
+output variables without actually computing them.
+It takes as input ``fgraph``, a :class:`FunctionGraph`; ``node``, a reference
+to the :class:`Op`'s :ref:`apply` node;
+and a list of :class:`Variables`\s (e.g. ``i0_shape``, ``i1_shape``, ...)
+which are the dimensions of the :class:`Op` input :ref:`variable`\s.
+:meth:`infer_shape` returns a list where each element is a tuple representing
+the shape of one output.
+This could be helpful if one only needs the shape of the output instead of the
+actual outputs, which can be useful, for instance, for rewriting
+procedures.
- def perform(self, node, inputs, output_storage):
- x = inputs[0]
- z = output_storage[0]
- z[0] = x * 2
+:meth:`L_op`
+^^^^^^^^^^^^^^^
- def infer_shape(self, fgraph, node, i0_shapes):
- return i0_shapes
+The :meth:`L_op` method is required if you want to differentiate some cost
+whose expression includes your :class:`Op`. The gradient is
+specified symbolically in this method. It takes three arguments ``inputs``, ``outputs`` and
+``output_gradients``, which are both lists of :ref:`variable`\s, and
+those must be operated on using PyTensor's symbolic language. The :meth:`L_op`
+method must return a list containing one :ref:`variable` for each
+input. Each returned :ref:`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 be defined to return a variable of type :class:`NullType`
+for that input. Likewise, if you have not implemented the gradient
+computation for some input, you may return a variable of type
+:class:`NullType` for that input. Please refer to :meth:`L_op` for a more detailed
+view.
+
+:meth:`R_op`
+^^^^^^^^^^^^^^^
+The :meth:`R_op` method is needed if you want :func:`pytensor.gradient.Rop` to
+work with your :class:`Op`.
- def grad(self, inputs, output_grads):
- return [output_grads[0] * 2]
+This function implements the application of the R-operator on the
+function represented by your :class:`Op`. Let's assume that function is :math:`f`,
+with input :math:`x`, applying the R-operator means computing the
+Jacobian of :math:`f` and right-multiplying it by :math:`v`, the evaluation
+point, namely: :math:`\frac{\partial f}{\partial x} v`.
- def R_op(self, inputs, eval_points):
- # R_op can receive None as eval_points.
- # That mean there is no diferientiable path through that input
- # If this imply that you cannot compute some outputs,
- # return None for those.
- if eval_points[0] is None:
- return eval_points
- return self.grad(inputs, eval_points)
- doubleOp1 = DoubleOp1()
+Example: :class:`Op` definition
+-------------------------------
- #Using itypes and otypes
+.. testcode:: example
+ import numpy as np
+ from pytensor.graph.op import Op
+ from pytensor.graph.basic import Apply, Variable
+ from pytensor.tensor import as_tensor_variable, TensorLike, TensorVariable
- class DoubleOp2(Op):
+ class DoubleOp1(Op):
__props__ = ()
- itypes = [pytensor.tensor.dmatrix]
- otypes = [pytensor.tensor.dmatrix]
+ def make_node(self, x: TensorLike) -> Apply:
+ # Convert (and require) x to be a TensorVariable
+ x = as_tensor_variable(x)
- def perform(self, node, inputs, output_storage):
+ # Validate input type
+ if not x.type.ndim == 2 and x.type.dtype == "float64":
+ raise TypeError("x must be a float64 matrix")
+
+ # Create an output variable of the same type as x
+ z = x.type()
+
+ # TensorVariables type include shape and dtype, so this is equivalent to the following
+ # z = pytensor.tensor.TensorType(dtype=x.type.dtype, shape=x.type.shape)()
+ # z = pytensor.tensor.tensor(dtype=x.type.dtype, shape=x.type.shape)
+ return Apply(self, [x], [z])
+
+ def perform(self, node: Apply, inputs: list[np.ndarray], output_storage: list[list[np.ndarray | None]]) -> None:
x = inputs[0]
z = output_storage[0]
+ # Numerical output based on numerical inputs (i.e., numpy arrays)
z[0] = x * 2
- def infer_shape(self, fgraph, node, i0_shapes):
- return i0_shapes
+ def infer_shape(self, fgraph: FunctionGraph, node: Apply, input_shapes: list[list[Variable]]) -> list[list[Variable]]:
+ # The output shape is the same as the input shape
+ return input_shapes
- def grad(self, inputs, output_grads):
+ def L_op(self, inputs: list[TensorVariable], outputs: list[TensorVariable], output_grads: list[TensorVariable]):
+ # Symbolic expression for the gradient
+ # For this Op, the inputs and outputs aren't part of the expression
+ # output_grads[0] is a TensorVariable!
return [output_grads[0] * 2]
- def R_op(self, inputs, eval_points):
+ def R_op(self, inputs: list[TensorVariable], eval_points: list[TensorVariable | None]) -> list[TensorVariable] | None:
# R_op can receive None as eval_points.
# That mean there is no diferientiable path through that input
# If this imply that you cannot compute some outputs,
# return None for those.
if eval_points[0] is None:
- return eval_points
- return self.grad(inputs, eval_points)
+ return None
+ # For this Op, the R_op is the same as the L_op
+ outputs = self(inputs)
+ return self.L_op(inputs, outputs, eval_points)
- doubleOp2 = DoubleOp2()
+ doubleOp1 = DoubleOp1()
+
+At a high level, the code fragment declares a class (e.g., ``DoubleOp1``) and then creates one instance of that class (e.g., ``doubleOp1``).
+
+As you'll see below, you can then pass an instantiated :ref:`variable`, such as ``x = tensor.matrix("x")`` to the instantiated :class:`Op`,
+to define a new :ref:`variable` that represents the output of applying the :class:`Op` to the input variable.
-At a high level, the code fragment declares a class (e.g., ``DoubleOp1``) and then
-creates one instance of it (e.g., ``doubleOp1``).
-
-We often gloss over this distinction, but will be precise here:
-``doubleOp1`` (the instance) is an :class:`Op`, not ``DoubleOp1`` (the class which is a
-subclass of :class:`Op`). You can call ``doubleOp1(tensor.vector())`` on a
-``Variable`` to build an expression, and in the expression there will be
-a ``.op`` attribute that refers to ``doubleOp1``.
-
-.. The first two methods in the :class:`Op` are relatively boilerplate: ``__eq__``
-.. and ``__hash__``.
-.. When two :class:`Op`\s are equal, PyTensor will merge their outputs if they are applied to the same inputs.
-.. The base class says two objects are equal if (and only if)
-.. they are the same object.
-.. Writing these boilerplate definitions ensures that the logic of the equality comparison is always explicit.
-
-.. It is an essential part of the :ref:`op_contract` that if two :class:`Op`\s compare
-.. equal, then they must compute the same result when presented with the same
-.. inputs. Here, if we allocated another instance of ``Fibby`` by typing ``fibby2
-.. = Fibby()`` then we would have two :class:`Op`\s that behave identically.
-..
-.. When should the implementation of ``__eq__`` be more complicated?
-.. If ``Fibby.__init__`` had parameters, then we could
-.. have configured ``fibby2`` differently from ``fibby`` by passing different
-.. arguments to the constructor. If we had done that, and if that different
-.. configuration made ``fibby2`` compute different results from ``fibby`` (for the
-.. same inputs) then we would have to add logic to the ``__eq__`` and ``__hash__``
-.. function so that he two ``Fibby`` :class:`Op`\s would *not be equal*. The reason why: PyTensor's merge
-.. optimization looks for :class:`Op`\s comparing equal and merges them. If two :class:`Op`\s compare
-.. equal but don't always produce equal results from equal inputs, then you might
-.. see wrong calculation.
-
-The ``make_node`` method creates a node to be included in the expression graph.
-It runs when we apply our :class:`Op` (``doubleOp1``) to the ``Variable`` (``x``), as
-in ``doubleOp1(tensor.vector())``.
-When an :class:`Op` has multiple inputs, their order in the inputs argument to ``Apply``
-is important: PyTensor will call ``make_node(*inputs)`` to copy the graph,
-so it is important not to change the semantics of the expression by changing
-the argument order.
-
-All the ``inputs`` and ``outputs`` arguments to :class:`Apply` must be :class:`Variable`\s.
+Under the hood, the :meth:`__call__` will call :meth:`make_node` method and then returns the output variable(s)
+of the :ref:`apply` that is returned by the method.
+
+The number and order of the inputs argument in the returned :ref:`apply` should match those in the :meth:`make_node`.
+PyTensor may decide to call :meth:`make_node` itself later to copy the graph or perform a generic rewrite.
+
+All the ``inputs`` and ``outputs`` arguments to the returned :ref:`apply` must be :ref:`variable`\s.
A common and easy way to ensure inputs are variables is to run them through
-``as_tensor_variable``. This function leaves :class:`TensorType` variables alone, raises
-an error for non-:class:`TensorType` variables, and copies any ``numpy.ndarray`` into
-the storage for a :class:`TensorType` :class:`Constant`. The :func:`make_node` method dictates the
-appropriate :class:`Type` for all output variables.
+``as_tensor_variable``. This function leaves :class:`TensorVariable` variables alone, raises
+an error for variables with an incompatible type, and copies any ``numpy.ndarray`` into
+the storage for a :class:`TensorConstant`.
-The :func:`perform` method implements the :class:`Op`'s mathematical logic in Python.
-The inputs (here ``x``) are passed by value, but a single output is returned
-indirectly as the first element of single-element lists. If ``doubleOp1`` had
-a second output, it would be stored in ``output_storage[1][0]``.
+The :meth:`perform` method implements the :class:`Op`'s mathematical logic in Python.
+The inputs (here ``x = inputs[0]``) are passed by value, and a single output is stored
+as the first element of a single-element list (here ``z = output_storage[0]``).
+If ``doubleOp1`` had a second output, it should be stored in ``output_storage[1][0]``.
In some execution modes, the output storage might contain the return value of
a previous call. That old value can be reused to avoid memory re-allocation,
@@ -399,68 +302,76 @@ You can try the new :class:`Op` as follows:
.. testcode:: example
- import numpy as np
- import pytensor
+ from pytensor import function
+ from pytensor.tensor import matrix
- x = pytensor.tensor.matrix()
- f = pytensor.function([x], DoubleOp1()(x))
- inp = np.random.random_sample((5, 4))
- out = f(inp)
- assert np.allclose(inp * 2, out)
- print(inp)
- print(out)
+ doubleOp1 = DoubleOp1()
-.. testoutput:: example
- :hide:
- :options: +ELLIPSIS, +SKIP
+ x = matrix("x")
+ out = doubleOp1(x)
+ assert out.type == x.type
- <BLANKLINE>
+ fn = function([x], out)
+ x_np = np.random.normal(size=(5, 4))
+ np.testing.assert_allclose(x_np * 2, fn(x_np))
-.. code-block:: none
- [[ 0.08257206 0.34308357 0.5288043 0.06582951]
- [ 0.65977826 0.10040307 0.5402353 0.55472296]
- [ 0.82358552 0.29502171 0.97387481 0.0080757 ]
- [ 0.77327215 0.65401857 0.76562992 0.94145702]
- [ 0.8452076 0.30500101 0.88430501 0.95818655]]
- [[ 0.16514411 0.68616713 1.0576086 0.13165902]
- [ 1.31955651 0.20080613 1.08047061 1.10944593]
- [ 1.64717104 0.59004341 1.94774962 0.0161514 ]
- [ 1.5465443 1.30803715 1.53125983 1.88291403]
- [ 1.6904152 0.61000201 1.76861002 1.9163731 ]]
+It's also a good idea to test the :meth:`infer_shape` implementation.
+To do this we can request a graph of the shape only:
-.. testcode:: example
+.. testcode::
- import numpy as np
- import pytensor
+ out_shape = out.shape
+ shape_fn = function([x], out_shape)
+ assert tuple(shape_fn(x_np)) == x_np.shape
- x = pytensor.tensor.matrix()
- f = pytensor.function([x], DoubleOp2()(x))
- inp = np.random.random_sample((5, 4))
- out = f(inp)
- assert np.allclose(inp * 2, out)
- print(inp)
- print(out)
+ # We can introspect the compiled function to confirm the Op is not evaluated
+ shape_fn.dprint()
+.. testoutput::
-.. testoutput:: example
- :hide:
- :options: +ELLIPSIS, +SKIP
+ MakeVector{dtype='int64'} [id A] 2
+ ├─ Shape_i{0} [id B] 1
+ │ └─ x [id C]
+ └─ Shape_i{1} [id D] 0
+ └─ x [id C]
- <BLANKLINE>
-.. code-block:: none
+Finally we should test the gradient implementation.
+For this we can use the ``pytensor.gradient.verify_grad`` utility which will compare the output of a gradient function with fininte differences.
+
+.. testcode::
+ from pytensor.gradient import verify_grad
+
+ rng = np.random.default_rng(42)
+ test_x = rng.normal(size=(5, 4))
+
+ # Raises if the gradient output is sufficiently different from the finite difference approximation.
+ verify_grad(doubleOp1, [test_x], rng=rng)
- [[ 0.02443785 0.67833979 0.91954769 0.95444365]
- [ 0.60853382 0.7770539 0.78163219 0.92838837]
- [ 0.04427765 0.37895602 0.23155797 0.4934699 ]
- [ 0.20551517 0.7419955 0.34500905 0.49347629]
- [ 0.24082769 0.49321452 0.24566545 0.15351132]]
- [[ 0.04887571 1.35667957 1.83909538 1.90888731]
- [ 1.21706764 1.55410779 1.56326439 1.85677674]
- [ 0.08855531 0.75791203 0.46311594 0.9869398 ]
- [ 0.41103034 1.48399101 0.69001811 0.98695258]
- [ 0.48165539 0.98642904 0.4913309 0.30702264]]
+
+Example: :attr:`itypes` and :attr:`otypes` definition
+-----------------------------------------------------
+
+Since the `Op` has a very strict type signature, we can use :attr:`itypes` and :attr:`otypes` instead of :meth:`make_node`:
+
+.. testcode:: example with itypes and otypes
+
+ from pytensor.tensor import dmatrix
+
+ class DoubleOp2(Op):
+ __props__ = ()
+
+ # inputs and output types must be float64 matrices
+ itypes = [dmatrix]
+ otypes = [dmatrix]
+
+ def perform(self, node, inputs, output_storage):
+ x = inputs[0]
+ z = output_storage[0]
+ z[0] = x * 2
+
+ doubleOp2 = DoubleOp2()
Example: :attr:`__props__` definition
@@ -470,15 +381,13 @@ We can modify the previous piece of code in order to demonstrate
the usage of the :attr:`__props__` attribute.
We create an :class:`Op` that takes a variable ``x`` and returns ``a*x+b``.
-We want to say that two such :class:`Op`\s are equal when their values of ``a``
-and ``b`` are equal.
+We want to say that two such :class:`Op`\s are equal when their values of ``a`` and ``b`` are equal.
.. testcode:: properties
- import pytensor
from pytensor.graph.op import Op
from pytensor.graph.basic import Apply
-
+ from pytensor.tensor import as_tensor_variable
class AXPBOp(Op):
"""
@@ -492,7 +401,7 @@ and ``b`` are equal.
super().__init__()
def make_node(self, x):
- x = pytensor.tensor.as_tensor_variable(x)
+ x = as_tensor_variable(x)
return Apply(self, [x], [x.type()])
def perform(self, node, inputs, output_storage):
@@ -500,22 +409,18 @@ and ``b`` are equal.
z = output_storage[0]
z[0] = self.a * x + self.b
- def infer_shape(self, fgraph, node, i0_shapes):
- return i0_shapes
-
- def grad(self, inputs, output_grads):
- return [self.a * output_grads[0]]
-
-The use of :attr:`__props__` saves
-the user the trouble of implementing :func:`__eq__` and :func:`__hash__`
-manually. It also generates a default :func:`__str__` method that prints the
-attribute names and their values.
+The use of :attr:`__props__` saves the user the trouble of implementing :meth:`__eq__` and :meth:`__hash__` manually.
+It also generates default :meth:`__repr__` and :meth:`__str__` methods that prints the attribute names and their values.
We can test this by running the following segment:
.. testcode:: properties
+ import numpy as np
+ from pytensor.tensor import matrix
+ from pytensor import function
+
mult4plus5op = AXPBOp(4, 5)
another_mult4plus5op = AXPBOp(4, 5)
mult2plus3op = AXPBOp(2, 3)
@@ -523,111 +428,317 @@ We can test this by running the following segment:
assert mult4plus5op == another_mult4plus5op
assert mult4plus5op != mult2plus3op
- x = pytensor.tensor.matrix()
- f = pytensor.function([x], mult4plus5op(x))
- g = pytensor.function([x], mult2plus3op(x))
+ x = matrix("x", dtype="float32")
+ f = function([x], mult4plus5op(x))
+ g = function([x], mult2plus3op(x))
+
+ inp = np.random.normal(size=(5, 4)).astype("float32")
+ np.testing.assert_allclose(4 * inp + 5, f(inp))
+ np.testing.assert_allclose(2 * inp + 3, g(inp))
+
+
+To demonstrate the use of equality, we will define the following graph: ``mult4plus5op(x) + another_mult4plus5op(x) + mult3plus2op(x)``.
+And confirm PyTensor infers it can reuse the first term in place of the second ``another_mult4plus5op(x)``.
+
+.. testcode:: exploiting equality
+
+ from pytensor.graph import rewrite_graph
+
+ graph = mult4plus5op(x) + another_mult4plus5op(x) + mult2plus3op(x)
+ print("Before:")
+ graph.dprint()
- inp = np.random.random_sample((5, 4)).astype(np.float32)
- assert np.allclose(4 * inp + 5, f(inp))
- assert np.allclose(2 * inp + 3, g(inp))
+ print("\nAfter:")
+ rewritten_graph = rewrite_graph(graph)
+ rewritten_graph.dprint()
-How To Test it
---------------
+.. testoutput::
+ Before:
+ Add [id A]
+ ├─ Add [id B]
+ │ ├─ AXPBOp{a=4, b=5} [id C]
+ │ │ └─ x [id D]
+ │ └─ AXPBOp{a=4, b=5} [id E]
+ │ └─ x [id D]
+ └─ AXPBOp{a=2, b=3} [id F]
+ └─ x [id D]
+
+ After:
+ Add [id A]
+ ├─ AXPBOp{a=4, b=5} [id B]
+ │ └─ x [id C]
+ ├─ AXPBOp{a=4, b=5} [id B]
+ │ └─ ···
+ └─ AXPBOp{a=2, b=3} [id D]
+ └─ x [id C]
+
+Note how after rewriting, the same variable [id B] is used twice.
+Also the string representation of the `Op` shows the values of the properties.
+
+
+Example: More complex :class:`Op`
+---------------------------------
+
+As a final example, we will create a multi-output :class:`Op` that takes a matrix and a vector and returns the matrix transposed and the sum of the vector.
+
+Furthermore, this :class:`Op` will work with batched dimensions, meaning we can pass in a 3D tensor or a 2D tensor (or more) and it will work as expected.
+To achieve this behavior we cannot use `itypes` and `otypes` as those encode specific number of dimensions.
+Instead we will have to define the `make_node` method.
+
+We need to be careful in the :meth:`L_op` method, as one of output gradients may be disconnected from the cost, in which case we should ignore its contribution.
+If both outputs are disconnected PyTensor will not bother calling the :meth:`L_op` method, so we don't need to worry about that case.
+
+.. testcode::
+
+ import pytensor.tensor as pt
+
+ from pytensor.graph.op import Op
+ from pytensor.graph.basic import Apply
+ from pytensor.gradient import DisconnectedType
+
+ class TransposeAndSumOp(Op):
+ __props__ = ()
+
+ def make_node(self, x, y):
+ # Convert to TensorVariables (and fail if not possible)
+ x = pt.as_tensor_variable(x)
+ y = pt.as_tensor_variable(y)
+
+ # Validate inputs dimensions
+ if x.type.ndim < 2:
+ raise TypeError("x must be at least a matrix")
+ if y.type.ndim < 1:
+ raise TypeError("y must be at least a vector")
+
+ # Create output variables
+ out1_static_shape = (*x.type.shape[:-2], x.type.shape[-1], x.type.shape[-2])
+ out1_dtype = x.type.dtype
+ out1 = pt.tensor(dtype=out1_dtype, shape=out1_static_shape)
+
+ out2_static_shape = y.type.shape[:-1]
+ out2_dtype = "float64" # hard-coded regardless of the input
+ out2 = pt.tensor(dtype=out2_dtype, shape=out2_static_shape)
+
+ return Apply(self, [x, y], [out1, out2])
+
+ def perform(self, node, inputs, output_storage):
+ x, y = inputs
+ out_1, out_2 = output_storage
+ out_1[0] = np.swapaxes(x, -1, -2)
+ out_2[0] = y.sum(-1).astype("float64")
+
+ def infer_shape(self, fgraph, node, input_shapes):
+ x_shapes, y_shapes = input_shapes
+ out1_shape = (*x_shapes[:-2], x_shapes[-1], x_shapes[-2])
+ out2_shape = y_shapes[:-1]
+ return [out1_shape, out2_shape]
+
+ def L_op(self, inputs, outputs, output_grads):
+ x, y = inputs
+ out1_grad, out2_grad = output_grads
+
+ if isinstance(out1_grad.type, DisconnectedType):
+ x_grad = DisconnectedType()()
+ else:
+ # Transpose the last two dimensions of the output gradient
+ x_grad = pt.swapaxes(out1_grad, -1, -2)
+
+ if isinstance(out2_grad.type, DisconnectedType):
+ y_grad = DisconnectedType()()
+ else:
+ # Broadcast the output gradient to the same shape as y
+ y_grad = pt.broadcast_to(pt.expand_dims(out2_grad, -1), y.shape)
+
+ return [x_grad, y_grad]
+
+Let's test the `Op` evaluation:
+
+.. testcode::
+
+ import numpy as np
+ from pytensor import function
+
+ transpose_and_sum_op = TransposeAndSumOp()
+
+ x = pt.tensor("x", shape=(5, None, 3), dtype="float32")
+ y = pt.matrix("y", shape=(2, 1), dtype="float32")
+ x_np = np.random.normal(size=(5, 4, 3)).astype(np.float32)
+ y_np = np.random.normal(size=(2, 1)).astype(np.float32)
+
+ out1, out2 = transpose_and_sum_op(x, y)
+
+ # Test the output types
+ assert out1.type.shape == (5, 3, None)
+ assert out1.type.dtype == "float32"
+ assert out2.type.shape == (2,)
+ assert out2.type.dtype == "float64"
+
+ # Test the perform method
+ f = function([x, y], [out1, out2])
+ out1_np, out2_np = f(x_np, y_np)
+ np.testing.assert_allclose(out1_np, x_np.swapaxes(-1, -2))
+ np.testing.assert_allclose(out2_np, y_np.sum(-1))
+
+
+And the shape inference:
+
+.. testcode::
+
+ out1_shape = out1.shape
+ out2_shape = out2.shape
+ shape_fn = function([x, y], [out1_shape, out2_shape])
+
+ out1_shape_np, out2_shape_np = shape_fn(x_np, y_np)
+ assert tuple(out1_shape_np) == out1_np.shape
+ assert tuple(out2_shape_np) == out2_np.shape
+
+ # We can introspect the compiled function to confirm the Op is not needed
+ shape_fn.dprint()
+
+.. testoutput::
+
+ MakeVector{dtype='int64'} [id A] 1
+ ├─ 5 [id B]
+ ├─ 3 [id C]
+ └─ Shape_i{1} [id D] 0
+ └─ x [id E]
+ DeepCopyOp [id F] 2
+ └─ [2] [id G]
+
+
+Finally, the gradient expression:
+
+Again, we can use pytensor `verify_grad` function to test the gradient implementation.
+Due to the presence of multiple outputs we need to pass a `Callable` instead of the `Op` instance.
+There are different cases we want to test: when both or just one of the outputs is connected to the cost
+
+
+.. testcode::
+ import warnings
+ import numpy as np
+ from pytensor.gradient import verify_grad
+
+ transpose_and_sum_op = TransposeAndSumOp()
+
+ def both_outs_connected(x, y):
+ out1, out2 = transpose_and_sum_op(x, y)
+ return out1.sum() + out2.sum()
+
+ def only_out1_connected(x, y):
+ out1, _ = transpose_and_sum_op(x, y)
+ return out1.sum()
+
+ def only_out2_connected(x, y):
+ _, out2 = transpose_and_sum_op(x, y)
+ return out2.sum()
+
+ rng = np.random.default_rng(seed=37)
+ x_np = rng.random((5, 4, 3)).astype(np.float32)
+ y_np = rng.random((2, 1)).astype(np.float32)
+ verify_grad(both_outs_connected, [x_np, y_np], rng=rng)
+
+ # PyTensor will raise a warning about the disconnected gradient
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ verify_grad(only_out1_connected, [x_np, y_np], rng=rng)
+ verify_grad(only_out2_connected, [x_np, y_np], rng=rng)
+
+We are filtering a warning about DisconnectTypes being returned by the gradient method.
+PyTensor would like to know how the outputs of the `Op` are connected to the input, which could be done with `connection_pattern`
+This was omitted for brevity, since it's a rare edge-case.
+
+
+Developer testing utilities
+---------------------------
+
+PyTensor has some functionalities to test for a correct implementation of an :class:`Op` and it's many methods.
+
+We have already seen some user facing helpers, but there are also test classes for :class:`Op` implementations
+that are added to the codebase, to be used with ``pytest``.
+
+Here we mention those that can be used to test the implementation of:
+ :meth:`infer_shape`
+ :meth:`L_op`
+ :meth:`R_op`
-PyTensor has some functionalities to simplify testing. These help test the
-:meth:`Op.infer_shape`, :meth:`Op.grad` and :meth:`Op.R_op` methods. Put the following code
-in a file and execute it with the ``pytest`` program.
Basic Tests
^^^^^^^^^^^
-Basic tests are done by you just by using the :class:`Op` and checking that it
-returns the right answer. If you detect an error, you must raise an
-exception. You can use the ``assert`` keyword to automatically raise an
-`AssertionError`.
+Basic tests are done by you just by using the :class:`Op` and checking that it returns the right answer.
+If you detect an error, you must raise an exception.
+
+You can use the ``assert`` keyword to automatically raise an `AssertionError`, or utilities in `numpy.testing`.
.. testcode:: tests
import numpy as np
- import pytensor
- from tests import unittest_tools as utt
+ from pytensor import function
+ from pytensor.tensor import matrix
+ from tests.unittest_tools import InferShapeTester
- class TestDouble(utt.InferShapeTester):
+ class TestDouble(InferShapeTester):
def setup_method(self):
super().setup_method()
self.op_class = DoubleOp
self.op = DoubleOp()
def test_basic(self):
- rng = np.random.default_rng(utt.fetch_seed())
+ rng = np.random.default_rng(377)
- x = pytensor.tensor.matrix()
+ x = matrix(x, dtype="float64")
f = pytensor.function([x], self.op(x))
- inp = np.asarray(rng.random((5, 4)), dtype=pytensor.config.floatX)
+ inp = np.asarray(rng.random((5, 4)), dtype="float64")
out = f(inp)
+
# Compare the result computed to the expected value.
- utt.assert_allclose(inp * 2, out)
+ np.testing.assert_allclose(inp * 2, out)
-We call ``utt.assert_allclose(expected_value, value)`` to compare
-NumPy ndarray.This raise an error message with more information. Also,
-the default tolerance can be changed with the PyTensor flags
-``config.tensor__cmp_sloppy`` that take values in 0, 1 and 2. The
-default value do the most strict comparison, 1 and 2 make less strict
-comparison.
-Testing the :meth:`Op.infer_shape`
+Testing the :meth:`infer_shape`
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-When a class inherits from the :class:`InferShapeTester` class, it gets the
-:meth:`InferShapeTester._compile_and_check` method that tests the :meth:`Op.infer_shape`
-method. It tests that the :class:`Op` gets rewritten out of the graph if only
-the shape of the output is needed and not the output
-itself. Additionally, it checks that the rewritten graph computes
-the correct shape, by comparing it to the actual shape of the computed
-output.
-
-:meth:`InferShapeTester._compile_and_check` compiles an PyTensor function. It takes as
-parameters the lists of input and output PyTensor variables, as would be
-provided to :func:`pytensor.function`, and a list of real values to pass to the
-compiled function. It also takes the :class:`Op` class as a parameter
-in order to verify that no instance of it appears in the shape-optimized graph.
-
-If there is an error, the function raises an exception. If you want to
-see it fail, you can implement an incorrect :meth:`Op.infer_shape`.
-
-When testing with input values with shapes that take the same value
-over different dimensions (for instance, a square matrix, or a ``tensor3``
-with shape ``(n, n, n)``, or ``(m, n, m)``), it is not possible to detect if
-the output shape was computed correctly, or if some shapes with the
-same value have been mixed up. For instance, if the :meth:`Op.infer_shape` uses
-the width of a matrix instead of its height, then testing with only
-square matrices will not detect the problem. This is why the
-:meth:`InferShapeTester._compile_and_check` method prints a warning in such a case. If
-your :class:`Op` works only with such matrices, you can disable the warning with the
-``warn=False`` parameter.
+When a class inherits from the :class:`InferShapeTester` class,
+it gets the :meth:`InferShapeTester._compile_and_check` method that tests the :meth:`infer_shape` method.
+It tests that the :class:`Op` gets rewritten out of the graph if only the shape of the output is needed and not the output itself.
+Additionally, it checks that the rewritten graph computes the correct shape, by comparing it to the actual shape of the computed output.
+
+:meth:`InferShapeTester._compile_and_check` compiles an PyTensor function.
+It takes as parameters the lists of input and output PyTensor variables,
+as would be provided to :func:`pytensor.function`,
+and a list of real values to pass to the compiled function.
+It also takes the :class:`Op` class as a parameter in order to verify that no instance of it appears in the shape-optimized graph.
+
+If there is an error, the function raises an exception.
+If you want to see it fail, you can implement an incorrect :meth:`infer_shape`.
+
+When testing with input values with shapes that take the same value over different dimensions
+(for instance, a square matrix, or a ``tensor3`` with shape ``(n, n, n)``, or ``(m, n, m)``),
+it is not possible to detect if the output shape was computed correctly,
+or if some shapes with the same value have been mixed up.
+For instance, if the :meth:`infer_shape` uses the width of a matrix instead of its height,
+then testing with only square matrices will not detect the problem.
+To avoid this the :meth:`InferShapeTester._compile_and_check` method prints a warning in such a case.
+If your :class:`Op` works only with such matrices, you can disable the warning with the ``warn=False`` parameter.
.. testcode:: tests
- from pytensor.configdefaults import config
- from tests import unittest_tools as utt
-
- class TestDouble(utt.InferShapeTester):
+ class TestDouble(InferShapeTester):
# [...] as previous tests.
def test_infer_shape(self):
- rng = np.random.default_rng(utt.fetch_seed())
- x = pytensor.tensor.matrix()
+ rng = np.random.default_rng(42)
+ x = matrix("x", dtype="float64")
self._compile_and_check(
[x], # pytensor.function inputs
[self.op(x)], # pytensor.function outputs
- # Always use not square matrix!
- # inputs data
- [np.asarray(rng.random((5, 4)), dtype=config.floatX)],
+ # Non-square inputs
+ [rng.random(size=(5, 4))],
# Op that should be removed from the graph.
self.op_class,
)
@@ -635,75 +746,52 @@ your :class:`Op` works only with such matrices, you can disable the warning with
Testing the gradient
^^^^^^^^^^^^^^^^^^^^
-The function :ref:`verify_grad <validating_grad>`
-verifies the gradient of an :class:`Op` or PyTensor graph. It compares the
-analytic (symbolically computed) gradient and the numeric
-gradient (computed through the Finite Difference Method).
+As shown above, the function :ref:`verify_grad <validating_grad>` verifies the gradient of an :class:`Op` or PyTensor graph.
+It compares the analytic (symbolically computed) gradient and the numeric gradient (computed through the Finite Difference Method).
-If there is an error, the function raises an exception. If you want to
-see it fail, you can implement an incorrect gradient (for instance, by removing
-the multiplication by 2).
+If there is an error, the function raises an exception.
+If you want to see it fail, you can implement an incorrect gradient
+(for instance, by removing the multiplication by 2).
.. testcode:: tests
def test_grad(self):
- rng = np.random.default_rng(utt.fetch_seed())
- tests.unittest_tools.verify_grad(
+ rng = np.random.default_rng(2024)
+ verify_grad(
self.op,
- [rng.random((5, 7, 2))]
+ [rng.random(size=(5, 7, 2))],
+ rng = rng,
)
Testing the Rop
^^^^^^^^^^^^^^^
-.. TODO: repair defective links in the following paragraph
-
-The class :class:`RopLop_checker` defines the functions
-:func:`RopLop_checker.check_mat_rop_lop`, :func:`RopLop_checker.check_rop_lop` and
-:func:`RopLop_checker.check_nondiff_rop`. These allow to test the
-implementation of the :meth:`Rop` method of a particular :class:`Op`.
+The class :class:`RopLopChecker` defines the methods
+:meth:`RopLopChecker.check_mat_rop_lop`, :meth:`RopLopChecker.check_rop_lop` and :meth:`RopLopChecker.check_nondiff_rop`.
+These allow to test the implementation of the :meth:`R_op` method of a particular :class:`Op`.
-For instance, to verify the :meth:`Rop` method of the ``DoubleOp``, you can use this:
+For instance, to verify the :meth:`R_op` method of the ``DoubleOp``, you can use this:
.. testcode:: tests
import numpy
import tests
- from tests.test_rop import RopLop_checker
- class TestDoubleRop(RopLop_checker):
+ from tests.test_rop import RopLopChecker
+
+ class TestDoubleRop(RopLopChecker):
+
def setUp(self):
super(TestDoubleRop, self).setUp()
+
def test_double_rop(self):
self.check_rop_lop(DoubleRop()(self.x), self.in_shape)
+
Running Your Tests
^^^^^^^^^^^^^^^^^^
To perform your tests, simply run ``pytest``.
-In-file
-"""""""
-
-One may also add a block of code similar to the following at the end
-of the file containing a specific test of interest and run the
-file. In this example, the test ``TestDoubleRop`` in the class
-``test_double_op`` would be performed.
-
-.. testcode:: tests
-
- if __name__ == '__main__':
- t = TestDoubleRop("test_double_rop")
- t.setUp()
- t.test_double_rop()
-
-We recommend that when we execute a file, we run all tests in that
-file. This can be done by adding this at the end of your test files:
-
-.. testcode:: tests
-
- if __name__ == '__main__':
- unittest.main()
-
Exercise
""""""""
@@ -713,41 +801,20 @@ Modify and execute to compute: ``x * y``.
Modify and execute the example to return two outputs: ``x + y`` and `jx - yj`.
-You can omit the :meth:`Rop` functions. Try to implement the testing apparatus
-described above.
-
-(Notice that PyTensor's current *elemwise fusion* rewrite is
-only applicable to computations involving a single output. Hence, to gain
-efficiency over the basic solution that is asked here, the two operations would
-have to be jointly rewritten explicitly in the code.)
-
-Random numbers in tests
-"""""""""""""""""""""""
-
-Making tests errors more reproducible is a good practice. To make
-tests more reproducible, one needs a way to get the same random
-numbers. This can be done by seeding NumPy's random number
-generator.
-
-For convenience, the classes :class:`InferShapeTester` and :class:`RopLop_checker`
-already do this for you. If you implement your own :meth:`setUp` method,
-don't forget to call the parent :meth:`setUp` method.
-
+You can omit the :meth:`Rop` functions. Try to implement the testing apparatus described above.
:download:`Solution<extending_pytensor_solution_1.py>`
:func:`as_op`
----------------------
+-------------
:func:`as_op` is a Python decorator that converts a Python function into a
basic PyTensor :class:`Op` that will call the supplied function during execution.
-This isn't the recommended way to build an :class:`Op`, but allows for a quick
-implementation.
+This isn't the recommended way to build an :class:`Op`, but allows for a quick implementation.
-It takes an optional :meth:`Op.infer_shape` parameter that must have this
-signature:
+It takes an optional :meth:`infer_shape` parameter that must have this signature:
.. code-block:: none
@@ -761,25 +828,24 @@ signature:
.. warning::
- Not providing a :meth:`Op.infer_shape` prevents shape-related
- rewrites from working with this :class:`Op`. For example
- ``your_op(inputs, ...).shape`` will need the :class:`Op` to be executed just
- to get the shape.
+ Not providing a :meth:`infer_shape` prevents shape-related rewrites from working with this :class:`Op`.
+ For example ``your_op(inputs, ...).shape`` will need the :class:`Op` to be executed just to get the shape.
.. note::
- As no grad is defined, this means you won't be able to
+ As no L_op is defined, this means you won't be able to
differentiate paths that include this :class:`Op`.
.. note::
- It converts the Python function to a callable object that takes as
+ It converts the Python function to a `Callable` object that takes as
inputs PyTensor variables that were declared.
.. note::
The python function wrapped by the :func:`as_op` decorator needs to return a new
data allocation, no views or in place modification of the input.
+
:func:`as_op` Example
^^^^^^^^^^^^^^^^^^^^^
@@ -791,14 +857,16 @@ signature:
from pytensor import function
from pytensor.compile.ops import as_op
-
def infer_shape_numpy_dot(fgraph, node, input_shapes):
ashp, bshp = input_shapes
return [ashp[:-1] + bshp[-1:]]
- @as_op(itypes=[pt.matrix, pt.matrix],
- otypes=[pt.matrix], infer_shape=infer_shape_numpy_dot)
+ @as_op(
+ itypes=[pt.dmatrix, pt.dmatrix],
+ otypes=[pt.dmatrix],
+ infer_shape=infer_shape_numpy_dot,
+ )
def numpy_dot(a, b):
return np.dot(a, b)
@@ -814,41 +882,32 @@ You can try it as follows:
out = f(inp1, inp2)
-.. _Documentation:
-
-Documentation and Coding Style
-------------------------------
-Please always respect the :ref:`quality_contributions` or your contribution
-will not be accepted.
+Final Note
+----------
-:class:`NanGuardMode` and :class:`AllocEmpty`
----------------------------------------------
+The section :ref:`Other Ops <other_ops>` includes more instructions for the following specific cases:
-:class:`NanGuardMode` help users find where in the graph ``NaN`` appear. But
-sometimes, we want some variables to not be checked. For example, in
-the old GPU back-end, we used a float32 :class:`CudaNdarray` to store the MRG
-random number generator state (they are integers). So if :class:`NanGuardMode`
-checked it, it would generate a false positive. Another case is related to
-:class:`AllocEmpty` or some computations on it (like done by :class:`Scan`).
+ - :ref:`scalar_ops`
+ - :ref:`sparse_ops`
+ - :ref:`openmp_ops`
-You can tell :class:`NanGuardMode` to do not check a variable with:
-:attr:`variable.tag.nan_guard_mode_check`. Also, this tag automatically
-follows that variable during rewriting. This mean if you tag a
-variable that get replaced by an inplace version, it will keep that
-tag.
+For defining C-based :class:`COp` see :ref:`creating_a_c_op`.
+For defining implementations for other backends see :ref:`creating_a_numba_jax_op`.
-Final Note
-----------
+.. note::
-A more extensive discussion of this section's content may be found in
-the advanced tutorial :ref:`Extending PyTensor<extending>`.
+ This is an introductory tutorial and as such it does not cover how to make
+ an :class:`Op` that returns a view or modifies the values in its inputs. Thus, all
+ :class:`Op`\s created with the instructions described here MUST return newly
+ allocated memory or reuse the memory provided in the parameter
+ ``output_storage`` of the :meth:`perform` method. See
+ :ref:`views_and_inplace` for an explanation on how to do this.
-The section :ref:`Other Ops <other_ops>` includes more instructions for
-the following specific cases:
+ If your :class:`Op` returns a view or changes the value of its inputs
+ without doing as prescribed in that page, PyTensor will run, but will
+ return correct results for some graphs and wrong results for others.
- - :ref:`scalar_ops`
- - :ref:`sparse_ops`
- - :ref:`Random ops <random_ops>`
- - :ref:`openmp_ops`
- - :ref:`numba_ops`
+ It is recommended that you run your tests in :class:`DebugMode`, since it
+ can help verify whether or not your :class:`Op` behaves correctly in this
+ regard.
diff --git a/doc/extending/other_ops.rst b/doc/extending/other_ops.rst
index fbad2ba48e..6bfc66b341 100644
--- a/doc/extending/other_ops.rst
+++ b/doc/extending/other_ops.rst
@@ -22,14 +22,6 @@ elemwise implementation will automatically have C code too. This
will enable the fusion of elemwise operations using your new scalar
operation. It is similar for reduction operations.
-Be careful about some possible problems in the definition of the
-``grad`` method, and about dependencies that may not be available. In
-particular, see the following fixes:
-`Fix to grad() methods
-<https://github.com/Theano/Theano/commit/002872ad97919b97eaf58e095044e3c3067668e4>`_
-and `impl() methods related to SciPy
-<https://github.com/Theano/Theano/commit/08d16c0aa6681fc53d8d0f40342551eb47ff536e>`_.
-
.. _sparse_ops:
Sparse Ops
@@ -116,43 +108,6 @@ needed sparse variable and data, you can use
many parameters, including parameters for the format (csr or csc), the shape, the
dtype, whether to have explicit 0 and whether to have unsorted indices.
-.. _random_ops:
-
-Random distribution
-===================
-
-We have 3 base random number generators. One that wraps NumPy's random
-generator, one that implements MRG31k3p and one that wraps CURAND.
-
-The recommended and 2nd faster is MRG. It works on the CPU and
-has more implemented distributions.
-
-The slowest is our wrapper on NumPy's random generator.
-
-We explain and provide advice on 3 possibles implementations of new
-distributions here:
-
-1. Extend our wrapper around NumPy random functions.
- See this `PR <https://github.com/Theano/Theano/pull/1607>`_ as an example.
-
-2. Extend MRG implementation by reusing existing PyTensor Op. Look into
- the ``PyTensor/sandbox/rng_mrg.py`` file and grep for all code about
- binomial(). This distribution uses the output of the uniform
- distribution and converts it to a binomial distribution with
- existing PyTensor operations. The tests go in
- ``PyTensor/sandbox/test_rng_mrg.py``
-
-3. Extend MRG implementation with a new Op that takes a uniform sample as
- input. Look in the ``PyTensor/sandbox/{rng_mrg,multinomial}.py`` file
- and its test in ``PyTensor/sandbox/test_multinomal.py``. This is
- recommended when current PyTensor ops aren't well suited to modify
- the uniform to the target distribution. This can happen in
- particular if there is a loop or complicated condition.
-
-.. note::
-
- In all cases, you must reuse the same interface as NumPy for compatibility.
-
.. _openmp_ops:
@@ -188,14 +143,6 @@ current convention.
same inputs and they execute 2 ConvOp that only differ on the
OpenMP parameter, we want them to be merged.
-.. _numba_ops:
-
-Numba Ops
-=========
-
-Want C speed without writing C code for your new Op? You can use Numba
-to generate the C code for you! Here is an `example
-Op <https://gist.github.com/nouiz/5492778#file-theano_op-py>`_ doing that.
.. _alternate_pytensor_types: