Skip to content

P
pytensor
提交 e301a423 authored 作者: ricardoV94's avatar ricardoV94 提交者： Ricardo Vieira
浏览文件
操作

Improve creating_an_op.rst

Changes: 1. Remove references to c-code which apply to `COp` but not `Op` 2. Fix failing doctests 3. Improve explanation of `make_node` 4. Emphasize distinction between itypes/otypes and make-node 5. Show `L_op` instead of `grad` 6. Show how to test `L_op` and `infer_shape` implementation 7. Simplify explanation of `__props__` and illustrate in example. 8. Introduce more complex multi-output Op to drive these details home 9. Remove old references to numba/ random variable Ops
上级 d10f2459
隐藏空白字符变更
内嵌 并排
正在显示 包含 590 行增加 和 587 行删除
doc/extending/creating_an_op.rst
浏览文件 @ e301a423
...@@ -4,37 +4,15 @@ ...@@ -4,37 +4,15 @@
Creating a new :class:`Op`: Python implementation Creating a new :class:`Op`: Python implementation
================================================= =================================================
So suppose you have looked through the library documentation and you don't see You may have looked through the library documentation but don't see a function that does what you want.
a function that does what you want.
If you can implement something in terms of an existing :ref:`op`, you should do that. If you can implement something in terms of an existing :class:`Op`, you should do that.
Odds are your function that uses existing PyTensor expressions is short, A PyTensor function that builds upon existing expressions will be better optimized, automatic differentiable, and
has no bugs, and potentially profits from rewrites that have already been work seamlessly across different backends.
implemented.
However, if you cannot implement an :class:`Op` in terms of an existing :class:`Op`, you have to However, if you cannot implement an :class:`Op` in terms of an existing :class:`Op`, you have to write a new one.
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.
This page will show how to implement some simple Python-based :class:`Op` that perform operations on numpy arrays.
PyTensor Graphs refresher PyTensor Graphs refresher
------------------------- -------------------------
...@@ -45,12 +23,12 @@ PyTensor Graphs refresher ...@@ -45,12 +23,12 @@ PyTensor Graphs refresher
PyTensor represents symbolic mathematical computations as graphs. Those graphs PyTensor represents symbolic mathematical computations as graphs. Those graphs
are bi-partite graphs (graphs with two types of nodes), they are composed of are bi-partite graphs (graphs with two types of nodes), they are composed of
interconnected :ref:`apply` and :ref:`variable` nodes. 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 intermediary values. As such, inputs and outputs of a graph are lists of PyTensor
:class:`Variable` nodes. :class:`Apply` nodes perform computation on these :ref:`variable` nodes. :ref:`apply` nodes perform computation on these
variables to produce new variables. Each :class:`Apply` node has a link to an 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 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 :ref:`graphstructures` for a more detailed explanation about the graph
structure. structure.
...@@ -58,338 +36,263 @@ structure. ...@@ -58,338 +36,263 @@ structure.
:class:`Op`'s basic methods :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 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 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 possibilities you may encounter or need. For that refer to
:ref:`op_contract`. :ref:`Op contract <op_contract>`.
.. testcode:: python .. 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.op import Op
from pytensor.graph.type import Type
class MyOp(Op): class MyOp(Op):
# Properties attribute # Properties attribute
__props__ = () __props__ : tuple[Any, ...] = ()
#itypes and otypes attributes are # Constructor, usually used only to set Op properties
#compulsory if make_node method is not defined. def __init__(self, *args):
#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):
pass pass
# Python implementation: # itypes and otypes attributes are compulsory if make_node method is not defined.
def perform(self, node, inputs_storage, output_storage): # They're the type of input and output respectively
pass itypes: list[Type] | None = None
otypes: list[Type] | None = None
# Other type of implementation # make_node is compulsory if itypes and otypes are not defined
# C implementation: [see pytensor web site for other functions] # make_node is more flexible: output types can be determined
def c_code(self, node, inputs, outputs, sub): # based on the input types and Op properties.
def make_node(self, *inputs) -> Apply:
pass pass
# Other implementations: # Performs the numerical evaluation of Op in Python. Required.
def make_thunk(self, node, storage_map, _, _2, impl=None): def perform(self, node: Apply, inputs_storage: list[Any], output_storage: list[list[Any]]) -> None:
pass pass
# optional: # Defines the symbolic expression for the L-operator based on the input and output variables
check_input = True # and the output gradient variables. Optional.
def L_op(self, inputs: list[Variable], outputs: list[Variable], output_grads: list[Variable]) -> list[Variable]:
def __init__(self, *args):
pass 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 pass
def R_op(self, inputs, eval_points): # Defines the 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 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 pass
An :class:`Op` has to implement some methods defined in the the interface of 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 :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 the method :meth:`make_node` or :attr:`itypes`, :attr:`otypes`, and :meth:`perform`.
implementation methods, either :meth:`Op.perform`, :meth:`COp.c_code`
or :meth:`Op.make_thunk`. :meth:`make_node`
^^^^^^^^^^^^^^^^^^^^^^^^
: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: :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:
- 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 - Checks that the inputs can be converted to :ref:`variable`\s whose types are compatible with the current :class:`Op`.
input types, it must raises an exception (such as :class:`TypeError`). If the :class:`Op` cannot be applied on the provided input types, it must raise an exception (such as :class:`TypeError`).
- it operates on the :class:`Variable`\s found in - Creates new output :ref:`variable`\s of a suitable symbolic :class:`Type` to serve as the outputs of this :class:`Op`'s application.
``*inputs`` in PyTensor's symbolic language to infer the type of - Returns an :ref:`apply` instance with the input and output :ref:`variable`\s, and itself as the :class:`Op`.
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 If :meth:`make_node` is not defined, the :attr:`itypes` and :attr:`otypes` are used by the :class:`Op`'s
application. :meth:`make_node` method to implement the functionality method mentioned above.
- it creates an :class:`Apply` instance with the input and output :class:`Variable`, and
return the :class:`Apply` instance.
:meth:`perform`
^^^^^^^^^^^^^^^^^^
:meth:`Op.perform` method defines the Python implementation of an :class:`Op`. :meth:`perform` method defines the Python implementation of an :class:`Op`.
It takes several arguments: It takes several arguments:
- ``node`` is a reference to an Apply node which was previously - ``node`` is a reference to an :ref:`apply` node which was previously
obtained via the :meth:`Op.make_node` method. It is typically not obtained via the :meth:`make_node` method. It is typically not
used in a simple :class:`Op`, but it contains symbolic information that used in a simple :class:`Op`, but it contains symbolic information that
could be required by a complex :class:`Op`. could be required by a complex :class:`Op`.
- ``inputs`` is a list of references to data which can be operated on using - ``inputs`` is a list of references to data which can be operated on using
non-symbolic statements, (i.e., statements in Python, Numpy). non-symbolic statements, (i.e., statements in Python, Numpy).
- ``output_storage`` is a list of storage cells where the output - ``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`. 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 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) symbolic output.
contained in ``output_storage``. PyTensor may sometimes allow ``output_storage`` elements to persist
A function Mode may allow ``output_storage`` elements to persist between evaluations, or it may reset ``output_storage`` cells to
between evaluations, or it may reset ``output_storage`` cells to hold a value of ``None``. It can also pre-allocate some memory
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
for the :class:`Op` to use. This feature can allow ``perform`` to reuse memory between calls, for example. If there is something
memory between calls, for example. If there is something preallocated in the ``output_storage``, it will be of the correct
preallocated in the ``output_storage``, it will be of the good dtype, but can have the wrong shape and have any stride pattern.
dtype, but can have the wrong shape and have any stride pattern.
:meth:`perform` method must be determined by the inputs.
:meth:`Op.perform` method must be determined by the inputs. That is to say, That is to say, when applied to identical inputs the method must return the same outputs.
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.
:class:`Op`'s auxiliary methods :class:`Op`'s auxiliary methods
------------------------------- -------------------------------
There are other methods that can be optionally defined by the :class:`Op`: There are other methods that can be optionally defined by the :class:`Op`:
:meth:`Op.__eq__` and :meth:`Op.__hash__` define respectively equality :attr:`__props__`
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.
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 :attr:`__props__` will also generate a suitable :meth:`__repr__` and :meth:`__str__` for your :class:`Op`.
from pytensor.graph.op import Op
from pytensor.graph.basic import Apply
class DoubleOp1(Op): :meth:`infer_shape`
__props__ = () ^^^^^^^^^^^^^^^^^^^^^^
def make_node(self, x): The :meth:`infer_shape` method allows an :class:`Op` to infer the shape of its
x = pytensor.tensor.as_tensor_variable(x) output variables without actually computing them.
# Note: using x_.type() is dangerous, as it copies x's broadcasting It takes as input ``fgraph``, a :class:`FunctionGraph`; ``node``, a reference
# behaviour to the :class:`Op`'s :ref:`apply` node;
return Apply(self, [x], [x.type()]) 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): :meth:`L_op`
x = inputs[0] ^^^^^^^^^^^^^^^
z = output_storage[0]
z[0] = x * 2
def infer_shape(self, fgraph, node, i0_shapes): The :meth:`L_op` method is required if you want to differentiate some cost
return i0_shapes 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): This function implements the application of the R-operator on the
return [output_grads[0] * 2] 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__ = () __props__ = ()
itypes = [pytensor.tensor.dmatrix] def make_node(self, x: TensorLike) -> Apply:
otypes = [pytensor.tensor.dmatrix] # 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] x = inputs[0]
z = output_storage[0] z = output_storage[0]
# Numerical output based on numerical inputs (i.e., numpy arrays)
z[0] = x * 2 z[0] = x * 2
def infer_shape(self, fgraph, node, i0_shapes): def infer_shape(self, fgraph: FunctionGraph, node: Apply, input_shapes: list[list[Variable]]) -> list[list[Variable]]:
return i0_shapes # 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] 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. # R_op can receive None as eval_points.
# That mean there is no diferientiable path through that input # That means there is no differentiable path through that input
# If this imply that you cannot compute some outputs, # If this imply that you cannot compute some outputs,
# return None for those. # return None for those.
if eval_points[0] is None: if eval_points[0] is None:
return eval_points return None
return self.grad(inputs, eval_points) # 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 Under the hood, the :meth:`__call__` will call :meth:`make_node` method and then returns the output variable(s)
creates one instance of it (e.g., ``doubleOp1``). of the :ref:`apply` that is returned by the method.
We often gloss over this distinction, but will be precise here: The number and order of the inputs argument in the returned :ref:`apply` should match those in the :meth:`make_node`.
``doubleOp1`` (the instance) is an :class:`Op`, not ``DoubleOp1`` (the class which is a PyTensor may decide to call :meth:`make_node` itself later to copy the graph or perform a generic rewrite.
subclass of :class:`Op`). You can call ``doubleOp1(tensor.vector())`` on a
``Variable`` to build an expression, and in the expression there will be All the ``inputs`` and ``outputs`` arguments to the returned :ref:`apply` must be :ref:`variable`\s.
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.
A common and easy way to ensure inputs are variables is to run them through 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 ``as_tensor_variable``. This function leaves :class:`TensorVariable` variables alone, raises
an error for non-:class:`TensorType` variables, and copies any ``numpy.ndarray`` into an error for variables with an incompatible type, and copies any ``numpy.ndarray`` into
the storage for a :class:`TensorType` :class:`Constant`. The :func:`make_node` method dictates the the storage for a :class:`TensorConstant`.
appropriate :class:`Type` for all output variables.
The :func:`perform` method implements the :class:`Op`'s mathematical logic in Python. The :meth:`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 The inputs (here ``x = inputs[0]``) are passed by value, and a single output is stored
indirectly as the first element of single-element lists. If ``doubleOp1`` had as the first element of a single-element list (here ``z = output_storage[0]``).
a second output, it would be stored in ``output_storage[1][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 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, 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: ...@@ -399,68 +302,76 @@ You can try the new :class:`Op` as follows:
.. testcode:: example .. testcode:: example
import numpy as np from pytensor import function
import pytensor from pytensor.tensor import matrix
x = pytensor.tensor.matrix() doubleOp1 = DoubleOp1()
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)
.. testoutput:: example x = matrix("x")
:hide: out = doubleOp1(x)
:options: +ELLIPSIS, +SKIP 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] It's also a good idea to test the :meth:`infer_shape` implementation.
[ 0.65977826 0.10040307 0.5402353 0.55472296] To do this we can request a graph of the shape only:
[ 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 ]]
.. testcode:: example .. testcode::
import numpy as np out_shape = out.shape
import pytensor shape_fn = function([x], out_shape)
assert tuple(shape_fn(x_np)) == x_np.shape
x = pytensor.tensor.matrix() # We can introspect the compiled function to confirm the Op is not evaluated
f = pytensor.function([x], DoubleOp2()(x)) shape_fn.dprint()
inp = np.random.random_sample((5, 4))
out = f(inp)
assert np.allclose(inp * 2, out)
print(inp)
print(out)
.. testoutput::
.. testoutput:: example MakeVector{dtype='int64'} [id A] 2
:hide: ├─ Shape_i{0} [id B] 1
:options: +ELLIPSIS, +SKIP │ └─ 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 finite 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] Example: :attr:`itypes` and :attr:`otypes` definition
[ 0.04427765 0.37895602 0.23155797 0.4934699 ] -----------------------------------------------------
[ 0.20551517 0.7419955 0.34500905 0.49347629]
[ 0.24082769 0.49321452 0.24566545 0.15351132]] Since the `Op` has a very strict type signature, we can use :attr:`itypes` and :attr:`otypes` instead of :meth:`make_node`:
[[ 0.04887571 1.35667957 1.83909538 1.90888731]
[ 1.21706764 1.55410779 1.56326439 1.85677674] .. testcode:: example with itypes and otypes
[ 0.08855531 0.75791203 0.46311594 0.9869398 ]
[ 0.41103034 1.48399101 0.69001811 0.98695258] from pytensor.tensor import dmatrix
[ 0.48165539 0.98642904 0.4913309 0.30702264]]
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 Example: :attr:`__props__` definition
...@@ -470,15 +381,13 @@ We can modify the previous piece of code in order to demonstrate ...@@ -470,15 +381,13 @@ We can modify the previous piece of code in order to demonstrate
the usage of the :attr:`__props__` attribute. the usage of the :attr:`__props__` attribute.
We create an :class:`Op` that takes a variable ``x`` and returns ``a*x+b``. 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`` We want to say that two such :class:`Op`\s are equal when their values of ``a`` and ``b`` are equal.
and ``b`` are equal.
.. testcode:: properties .. testcode:: properties
import pytensor
from pytensor.graph.op import Op from pytensor.graph.op import Op
from pytensor.graph.basic import Apply from pytensor.graph.basic import Apply
from pytensor.tensor import as_tensor_variable
class AXPBOp(Op): class AXPBOp(Op):
""" """
...@@ -492,7 +401,7 @@ and ``b`` are equal. ...@@ -492,7 +401,7 @@ and ``b`` are equal.
super().__init__() super().__init__()
def make_node(self, x): def make_node(self, x):
x = pytensor.tensor.as_tensor_variable(x) x = as_tensor_variable(x)
return Apply(self, [x], [x.type()]) return Apply(self, [x], [x.type()])
def perform(self, node, inputs, output_storage): def perform(self, node, inputs, output_storage):
...@@ -500,22 +409,18 @@ and ``b`` are equal. ...@@ -500,22 +409,18 @@ and ``b`` are equal.
z = output_storage[0] z = output_storage[0]
z[0] = self.a * x + self.b 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 :meth:`__eq__` and :meth:`__hash__` manually.
The use of :attr:`__props__` saves It also generates default :meth:`__repr__` and :meth:`__str__` methods that prints the attribute names and their values.
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.
We can test this by running the following segment: We can test this by running the following segment:
.. testcode:: properties .. testcode:: properties
import numpy as np
from pytensor.tensor import matrix
from pytensor import function
mult4plus5op = AXPBOp(4, 5) mult4plus5op = AXPBOp(4, 5)
another_mult4plus5op = AXPBOp(4, 5) another_mult4plus5op = AXPBOp(4, 5)
mult2plus3op = AXPBOp(2, 3) mult2plus3op = AXPBOp(2, 3)
...@@ -523,111 +428,317 @@ We can test this by running the following segment: ...@@ -523,111 +428,317 @@ We can test this by running the following segment:
assert mult4plus5op == another_mult4plus5op assert mult4plus5op == another_mult4plus5op
assert mult4plus5op != mult2plus3op assert mult4plus5op != mult2plus3op
x = pytensor.tensor.matrix() x = matrix("x", dtype="float32")
f = pytensor.function([x], mult4plus5op(x)) f = function([x], mult4plus5op(x))
g = pytensor.function([x], mult2plus3op(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) print("\nAfter:")
assert np.allclose(4 * inp + 5, f(inp)) rewritten_graph = rewrite_graph(graph)
assert np.allclose(2 * inp + 3, g(inp)) 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
^^^^^^^^^^^ ^^^^^^^^^^^
Basic tests are done by you just by using the :class:`Op` and checking that it Basic tests are done by you just by using the :class:`Op` and checking that it returns the right answer.
returns the right answer. If you detect an error, you must raise an If you detect an error, you must raise an exception.
exception. You can use the ``assert`` keyword to automatically raise an
`AssertionError`. You can use the ``assert`` keyword to automatically raise an `AssertionError`, or utilities in `numpy.testing`.
.. testcode:: tests .. testcode:: tests
import numpy as np import numpy as np
import pytensor from pytensor import function
from tests import unittest_tools as utt from pytensor.tensor import matrix
from tests.unittest_tools import InferShapeTester
class TestDouble(utt.InferShapeTester): class TestDouble(InferShapeTester):
def setup_method(self): def setup_method(self):
super().setup_method() super().setup_method()
self.op_class = DoubleOp self.op_class = DoubleOp
self.op = DoubleOp() self.op = DoubleOp()
def test_basic(self): 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)) 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) out = f(inp)
# Compare the result computed to the expected value. # 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 When a class inherits from the :class:`InferShapeTester` class,
:meth:`InferShapeTester._compile_and_check` method that tests the :meth:`Op.infer_shape` it gets the :meth:`InferShapeTester._compile_and_check` method that tests the :meth:`infer_shape` method.
method. It tests that the :class:`Op` gets rewritten out of the graph if only 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.
the shape of the output is needed and not the output Additionally, it checks that the rewritten graph computes the correct shape, by comparing it to the actual shape of the computed output.
itself. Additionally, it checks that the rewritten graph computes
the correct shape, by comparing it to the actual shape of the computed :meth:`InferShapeTester._compile_and_check` compiles an PyTensor function.
output. It takes as parameters the lists of input and output PyTensor variables,
as would be provided to :func:`pytensor.function`,
:meth:`InferShapeTester._compile_and_check` compiles an PyTensor function. It takes as and a list of real values to pass to the compiled function.
parameters the lists of input and output PyTensor variables, as would be 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.
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 If there is an error, the function raises an exception.
in order to verify that no instance of it appears in the shape-optimized graph. If you want to see it fail, you can implement an incorrect :meth:`infer_shape`.
If there is an error, the function raises an exception. If you want to When testing with input values with shapes that take the same value over different dimensions
see it fail, you can implement an incorrect :meth:`Op.infer_shape`. (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,
When testing with input values with shapes that take the same value or if some shapes with the same value have been mixed up.
over different dimensions (for instance, a square matrix, or a ``tensor3`` For instance, if the :meth:`infer_shape` uses the width of a matrix instead of its height,
with shape ``(n, n, n)``, or ``(m, n, m)``), it is not possible to detect if then testing with only square matrices will not detect the problem.
the output shape was computed correctly, or if some shapes with the To avoid this the :meth:`InferShapeTester._compile_and_check` method prints a warning in such a case.
same value have been mixed up. For instance, if the :meth:`Op.infer_shape` uses If your :class:`Op` works only with such matrices, you can disable the warning with the ``warn=False`` parameter.
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.
.. testcode:: tests .. testcode:: tests
from pytensor.configdefaults import config
from tests import unittest_tools as utt
class TestDouble(utt.InferShapeTester): class TestDouble(InferShapeTester):
# [...] as previous tests. # [...] as previous tests.
def test_infer_shape(self): def test_infer_shape(self):
rng = np.random.default_rng(utt.fetch_seed()) rng = np.random.default_rng(42)
x = pytensor.tensor.matrix() x = matrix("x", dtype="float64")
self._compile_and_check( self._compile_and_check(
[x], # pytensor.function inputs [x], # pytensor.function inputs
[self.op(x)], # pytensor.function outputs [self.op(x)], # pytensor.function outputs
# Always use not square matrix! # Non-square inputs
# inputs data [rng.random(size=(5, 4))],
[np.asarray(rng.random((5, 4)), dtype=config.floatX)],
# Op that should be removed from the graph. # Op that should be removed from the graph.
self.op_class, self.op_class,
) )
...@@ -635,75 +746,49 @@ your :class:`Op` works only with such matrices, you can disable the warning with ...@@ -635,75 +746,49 @@ your :class:`Op` works only with such matrices, you can disable the warning with
Testing the gradient Testing the gradient
^^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^
The function :ref:`verify_grad <validating_grad>` As shown above, the function :ref:`verify_grad <validating_grad>` verifies the gradient of an :class:`Op` or PyTensor graph.
verifies the gradient of an :class:`Op` or PyTensor graph. It compares the It compares the analytic (symbolically computed) gradient and the numeric gradient (computed through the Finite Difference Method).
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 If there is an error, the function raises an exception.
see it fail, you can implement an incorrect gradient (for instance, by removing If you want to see it fail, you can implement an incorrect gradient
the multiplication by 2). (for instance, by removing the multiplication by 2).
.. testcode:: tests .. testcode:: tests
def test_grad(self): def test_grad(self):
rng = np.random.default_rng(utt.fetch_seed()) rng = np.random.default_rng(2024)
tests.unittest_tools.verify_grad( verify_grad(
self.op, self.op,
[rng.random((5, 7, 2))] [rng.random(size=(5, 7, 2))],
rng = rng,
) )
Testing the Rop Testing the Rop
^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^
.. TODO: repair defective links in the following paragraph The class :class:`RopLopChecker` defines the methods
:meth:`RopLopChecker.check_mat_rop_lop`, :meth:`RopLopChecker.check_rop_lop` and :meth:`RopLopChecker.check_nondiff_rop`.
The class :class:`RopLop_checker` defines the functions These allow to test the implementation of the :meth:`R_op` method of a particular :class:`Op`.
: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`.
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 .. testcode:: tests
import numpy import numpy
import tests import tests
from tests.test_rop import RopLop_checker from tests.test_rop import RopLopChecker
class TestDoubleRop(RopLop_checker):
def setUp(self): class TestDoubleOpRop(RopLopChecker):
super(TestDoubleRop, self).setUp()
def test_double_rop(self): def test_double_rop(self):
self.check_rop_lop(DoubleRop()(self.x), self.in_shape) self.check_rop_lop(DoubleOp()(self.x), self.in_shape)
Running Your Tests Running Your Tests
^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^
To perform your tests, simply run ``pytest``. 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 Exercise
"""""""" """"""""
...@@ -713,41 +798,20 @@ Modify and execute to compute: ``x * y``. ...@@ -713,41 +798,20 @@ Modify and execute to compute: ``x * y``.
Modify and execute the example to return two outputs: ``x + y`` and `jx - yj`. 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 You can omit the :meth:`Rop` functions. Try to implement the testing apparatus described above.
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.
:download:`Solution<extending_pytensor_solution_1.py>` :download:`Solution<extending_pytensor_solution_1.py>`
:func:`as_op` :func:`as_op`
--------------------- -------------
:func:`as_op` is a Python decorator that converts a Python function into a :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. 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 This isn't the recommended way to build an :class:`Op`, but allows for a quick implementation.
implementation.
It takes an optional :meth:`Op.infer_shape` parameter that must have this It takes an optional :meth:`infer_shape` parameter that must have this signature:
signature:
.. code-block:: none .. code-block:: none
...@@ -761,25 +825,24 @@ signature: ...@@ -761,25 +825,24 @@ signature:
.. warning:: .. warning::
Not providing a :meth:`Op.infer_shape` prevents shape-related Not providing a :meth:`infer_shape` prevents shape-related rewrites from working with this :class:`Op`.
rewrites from working with this :class:`Op`. For example For example ``your_op(inputs, ...).shape`` will need the :class:`Op` to be executed just to get the shape.
``your_op(inputs, ...).shape`` will need the :class:`Op` to be executed just
to get the shape.
.. note:: .. 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`. differentiate paths that include this :class:`Op`.
.. note:: .. 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. inputs PyTensor variables that were declared.
.. note:: .. note::
The python function wrapped by the :func:`as_op` decorator needs to return a new 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. data allocation, no views or in place modification of the input.
:func:`as_op` Example :func:`as_op` Example
^^^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^^
...@@ -791,14 +854,16 @@ signature: ...@@ -791,14 +854,16 @@ signature:
from pytensor import function from pytensor import function
from pytensor.compile.ops import as_op from pytensor.compile.ops import as_op
def infer_shape_numpy_dot(fgraph, node, input_shapes): def infer_shape_numpy_dot(fgraph, node, input_shapes):
ashp, bshp = input_shapes ashp, bshp = input_shapes
return [ashp[:-1] + bshp[-1:]] return [ashp[:-1] + bshp[-1:]]
@as_op(itypes=[pt.matrix, pt.matrix], @as_op(
otypes=[pt.matrix], infer_shape=infer_shape_numpy_dot) itypes=[pt.dmatrix, pt.dmatrix],
otypes=[pt.dmatrix],
infer_shape=infer_shape_numpy_dot,
)
def numpy_dot(a, b): def numpy_dot(a, b):
return np.dot(a, b) return np.dot(a, b)
...@@ -814,41 +879,32 @@ You can try it as follows: ...@@ -814,41 +879,32 @@ You can try it as follows:
out = f(inp1, inp2) out = f(inp1, inp2)
.. _Documentation: Final Note
----------
Documentation and Coding Style
------------------------------
Please always respect the :ref:`quality_contributions` or your contribution
will not be accepted.
: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 - :ref:`scalar_ops`
sometimes, we want some variables to not be checked. For example, in - :ref:`sparse_ops`
the old GPU back-end, we used a float32 :class:`CudaNdarray` to store the MRG - :ref:`openmp_ops`
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`).
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 This is an introductory tutorial and as such it does not cover how to make
the advanced tutorial :ref:`Extending PyTensor<extending>`. 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 If your :class:`Op` returns a view or changes the value of its inputs
the following specific cases: 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` It is recommended that you run your tests in :class:`DebugMode`, since it
- :ref:`sparse_ops` can help verify whether or not your :class:`Op` behaves correctly in this
- :ref:`Random ops <random_ops>` regard.
- :ref:`openmp_ops`
- :ref:`numba_ops`
doc/extending/other_ops.rst
浏览文件 @ e301a423
...@@ -22,14 +22,6 @@ elemwise implementation will automatically have C code too. This ...@@ -22,14 +22,6 @@ elemwise implementation will automatically have C code too. This
will enable the fusion of elemwise operations using your new scalar will enable the fusion of elemwise operations using your new scalar
operation. It is similar for reduction operations. 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:
Sparse Ops Sparse Ops
...@@ -116,43 +108,6 @@ needed sparse variable and data, you can use ...@@ -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 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. 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: .. _openmp_ops:
...@@ -188,14 +143,6 @@ current convention. ...@@ -188,14 +143,6 @@ current convention.
same inputs and they execute 2 ConvOp that only differ on the same inputs and they execute 2 ConvOp that only differ on the
OpenMP parameter, we want them to be merged. 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: .. _alternate_pytensor_types:
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论