.. _op_contract: ============= :class:`Op`\s ============= An :class:`Op` is a :ref:`graph object ` that defines and performs computations in a graph. It has to define the following methods. .. function:: make_node(*inputs) This method is responsible for creating output :class:`Variable`\s of a suitable symbolic `Type` to serve as the outputs of this :Class:`Op`'s application. The :class:`Variable`\s found in ``*inputs`` must be operated on using PyTensor's symbolic language to compute the symbolic output :class:`Variable`\s. This method should put these outputs into an :class:`Apply` instance, and return the :class:`Apply` instance. This method creates an :class:`Apply` node representing the application of the `Op` on the inputs provided. If the `Op` cannot be applied to these inputs, it must raise an appropriate exception. The inputs of the :class:`Apply` instance returned by this call must be ordered correctly: a subsequent ``self.make_node(*apply.inputs)`` must produce something equivalent to the first ``apply``. .. function:: perform(node, inputs, output_storage) This method computes the function associated to this :class:`Op`. ``node`` is an :class:`Apply` node created by the :class:`Op`'s :meth:`Op.make_node` method. ``inputs`` is a list of references to data to operate on using non-symbolic statements, (i.e., statements in Python, NumPy). ``output_storage`` is a list of storage cells where the variables of the computation must be put. More specifically: - ``node``: This is a reference to an :class:`Apply` node which was previously obtained via the :meth:`Op.make_node` method. It is typically not used in simple :class:`Op`\s, but it contains symbolic information that could be required for complex :class:`Op`\s. - ``inputs``: This is a list of data from which the values stored in ``output_storage`` are to be computed using non-symbolic language. - ``output_storage``: This is a list of storage cells where the output is to be stored. A storage cell is a one-element list. It is forbidden to change the length of the list(s) contained in ``output_storage``. 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. This is a situation where the ``node`` argument can come in handy. A function :class:`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 :meth:`Op.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. This method must be determined by the inputs. That is to say, if it is evaluated once on inputs A and returned B, then if ever inputs C, equal to A, are presented again, then outputs equal to B must be returned again. You must be careful about aliasing outputs to inputs, and making modifications to any of the inputs. See :ref:`Views and inplace operations ` before writing a :meth:`Op.perform` implementation that does either of these things. .. function:: __eq__(other) ``other`` is also an :class:`Op`. Returning ``True`` here is a promise to the rewrite system that the other :class:`Op` will produce exactly the same graph effects (e.g. from its :meth:`Op.perform`) as this one, given identical inputs. This means it will produce the same output values, it will destroy the same inputs (same :attr:`Op.destroy_map`), and will alias outputs to the same inputs (same :attr:`Op.view_map`). For more details, see :ref:`views_and_inplace`. .. note:: If you set ``__props__``, this will be automatically generated. .. function:: __hash__() If two :class:`Op` instances compare equal, then they **must** return the same hash value. Equally important, this hash value must not change during the lifetime of self. :class:`Op` instances should be immutable in this sense. .. note:: If you set :attr:`Op.__props__`, this will be automatically generated. .. op_optional: Optional methods or attributes ============================== .. attribute:: __props__ Default: Undefined Must be a tuple. Lists the name of the attributes which influence the computation performed. This will also enable the automatic generation of appropriate ``__eq__``, ``__hash__`` and ``__str__`` methods. Should be set to ``()`` if you have no attributes that are relevant to the computation to generate the methods. .. versionadded:: 0.7 .. attribute:: default_output Default: None If this member variable is an integer, then the default implementation of ``__call__`` will return ``node.outputs[self.default_output]``, where ``node`` was returned by :meth:`Op.make_node`. Otherwise, the entire list of outputs will be returned, unless it is of length 1, where the single element will be returned by itself. .. function:: make_thunk(node, storage_map, compute_map, no_recycling, impl=None) This function must return a thunk, that is a zero-arguments function that encapsulates the computation to be performed by this :class:`Op` on the arguments of the node. :param node: :class:`Apply` instance The node for which a thunk is requested. :param storage_map: dict of lists This 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. :param compute_map: dict of lists This 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. :param no_recycling: WRITEME WRITEME :param impl: None, 'c' or 'py' Which implementation to use. The returned function must ensure that is sets the computed variables as computed in the `compute_map`. Defining this function removes the requirement for :meth:`perform` or C code, as you will define the thunk for the computation yourself. .. function:: __call__(*inputs, **kwargs) By default this is a convenience function which calls :meth:`make_node` with the supplied arguments and returns the result indexed by `default_output`. This can be overridden by subclasses to do anything else, but must return either an PyTensor :class:`Variable` or a list of :class:`Variable`\s. If you feel the need to override `__call__` to change the graph based on the arguments, you should instead create a function that will use your :class:`Op` and build the graphs that you want and call that instead of the :class:`Op` instance directly. .. function:: infer_shape(fgraph, node, shapes) This function is needed for shape rewrites. ``shapes`` is a list with one tuple for each input of the :class:`Apply` node (which corresponds to the inputs of the :class:`Op`). Each tuple contains as many elements as the number of dimensions of the corresponding input. The value of each element is the shape (number of items) along the corresponding dimension of that specific input. While this might sound complicated, it is nothing more than the shape of each input as symbolic variables (one per dimension). The function should return a list with one tuple for each output. Each tuple should contain the corresponding output's computed shape. Implementing this method will allow PyTensor to compute the output's shape without computing the output itself, potentially sparing you a costly recomputation. .. function:: flops(inputs, outputs) It is only used to have more information printed by the memory profiler. It makes it print the mega flops and giga flops per second for each apply node. It takes as inputs two lists: one for the inputs and one for the outputs. They contain tuples that are the shapes of the corresponding inputs/outputs. .. function:: __str__() This allows you to specify a more informative string representation of your :class:`Op`. If an `Op` has parameters, it is highly recommended to have the ``__str__`` method include the name of the :class:`Op` and the :Class:`Op`'s parameters' values. .. note:: If you set `__props__`, this will be automatically generated. You can still override it for custom output. .. function:: do_constant_folding(fgraph, node) Default: Return ``True`` By default when rewrites are enabled, we remove during function compilation :class:`Apply` nodes whose inputs are all constants. We replace the :class:`Apply` node with an PyTensor constant variable. This way, the :class:`Apply` node is not executed at each function call. If you want to force the execution of an :class:`Op` during the function call, make do_constant_folding return False. As done in the Alloc :class:`Op`, you can return False only in some cases by analyzing the graph from the node parameter. .. function:: debug_perform(node, inputs, output_storage) Undefined by default. If you define this function then it will be used instead of C code or :meth:`Op.perform` to do the computation while debugging (currently DebugMode, but others may also use it in the future). It has the same signature and contract as :meth:`Op.perform`. This enables :class:`Op`\s that cause trouble with DebugMode with their normal behaviour to adopt a different one when run under that mode. If your :class:`Op` doesn't have any problems, don't implement this. If you want your :class:`Op` to work with :func:`pytensor.gradient.grad` you also need to implement the functions described below. Gradient ======== These are the function required to work with :func:`pytensor.gradient.grad`. .. function:: grad(inputs, output_gradients) If the :class:`Op` being defined is differentiable, its gradient may be specified symbolically in this method. Both ``inputs`` and ``output_gradients`` are lists of symbolic PyTensor :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. :mod:`pytensor.gradient` contains convenience methods that can construct the variable for you: :func:`pytensor.gradient.grad_undefined` and :func:`pytensor.gradient.grad_not_implemented`, respectively. If an element of ``output_gradient`` is of type :class:`pytensor.gradient.DisconnectedType`, it means that the cost is not a function of this output. If any of the :class:`Op`'s inputs participate in the computation of only disconnected outputs, then :meth:`Op.grad` should return :class:`DisconnectedType` variables for those inputs. If the :meth:`Op.grad` method is not defined, then PyTensor assumes it has been forgotten. Symbolic differentiation will fail on a graph that includes this :class:`Op`. It must be understood that the :meth:`Op.grad` method is not meant to return the gradient of the :class:`Op`'s output. :func:`pytensor.grad` computes gradients; :meth:`Op.grad` is a helper function that computes terms that appear in gradients. If an :class:`Op` has a single vector-valued output ``y`` and a single vector-valued input ``x``, then the :meth:`Op.grad` method will be passed ``x`` and a second vector ``z``. Define ``J`` to be the Jacobian of ``y`` with respect to ``x``. The :meth:`Op.grad` method should return ``dot(J.T,z)``. When :func:`pytensor.grad` calls the :meth:`Op.grad` method, it will set ``z`` to be the gradient of the cost ``C`` with respect to ``y``. If this :class:`Op` is the only :class:`Op` that acts on ``x``, then ``dot(J.T,z)`` is the gradient of C with respect to ``x``. If there are other :class:`Op`\s that act on ``x``, :func:`pytensor.grad` will have to add up the terms of ``x``'s gradient contributed by the other :meth:`Op.grad` method. In practice, an :class:`Op`'s input and output are rarely implemented as single vectors. Even if an :class:`Op`'s output consists of a list containing a scalar, a sparse matrix, and a 4D tensor, you can think of these objects as being formed by rearranging a vector. Likewise for the input. In this view, the values computed by the :meth:`Op.grad` method still represent a Jacobian-vector product. In practice, it is probably not a good idea to explicitly construct the Jacobian, which might be very large and very sparse. However, the returned value should be equal to the Jacobian-vector product. So long as you implement this product correctly, you need not understand what :func:`pytensor.gradient.grad` is doing, but for the curious the mathematical justification is as follows: In essence, the :meth:`Op.grad` method must simply implement through symbolic :class:`Variable`\s and operations the chain rule of differential calculus. The chain rule is the mathematical procedure that allows one to calculate the total derivative :math:`\frac{d C}{d x}` of the final scalar symbolic `Variable` ``C`` with respect to a primitive symbolic :class:`Variable` x found in the list ``inputs``. The :meth:`Op.grad` method does this using ``output_gradients`` which provides the total derivative :math:`\frac{d C}{d f}` of ``C`` with respect to a symbolic :class:`Variable` that is returned by the `Op` (this is provided in ``output_gradients``), as well as the knowledge of the total derivative :math:`\frac{d f}{d x}` of the latter with respect to the primitive :class:`Variable` (this has to be computed). In mathematics, the total derivative of a scalar variable :math:`C` with respect to a vector of scalar variables :math:`x`, i.e. the gradient, is customarily represented as the row vector of the partial derivatives, whereas the total derivative of a vector of scalar variables :math:`f` with respect to another :math:`x`, is customarily represented by the matrix of the partial derivatives, i.e. the Jacobian matrix. In this convenient setting, the chain rule says that the gradient of the final scalar variable :math:`C` with respect to the primitive scalar variables in :math:`x` through those in :math:`f` is simply given by the matrix product: :math:`\frac{d C}{d x} = \frac{d C}{d f} * \frac{d f}{d x}`. Here, the chain rule must be implemented in a similar but slightly more complex setting: PyTensor provides in the list ``output_gradients`` one gradient for each of the :class:`Variable`\s returned by the `Op`. Where :math:`f` is one such particular :class:`Variable`, the corresponding gradient found in ``output_gradients`` and representing :math:`\frac{d C}{d f}` is provided with a shape similar to :math:`f` and thus not necessarily as a row vector of scalars. Furthermore, for each :class:`Variable` :math:`x` of the :class:`Op`'s list of input variables ``inputs``, the returned gradient representing :math:`\frac{d C}{d x}` must have a shape similar to that of :class:`Variable` x. If the output list of the :class:`Op` is :math:`[f_1, ... f_n]`, then the list ``output_gradients`` is :math:`[grad_{f_1}(C), grad_{f_2}(C), ... , grad_{f_n}(C)]`. If ``inputs`` consists of the list :math:`[x_1, ..., x_m]`, then `Op.grad` should return the list :math:`[grad_{x_1}(C), grad_{x_2}(C), ..., grad_{x_m}(C)]`, where :math:`(grad_{y}(Z))_i = \frac{\partial Z}{\partial y_i}` (and :math:`i` can stand for multiple dimensions). In other words, :meth:`Op.grad` does not return :math:`\frac{d f_i}{d x_j}`, but instead the appropriate dot product specified by the chain rule: :math:`\frac{d C}{d x_j} = \frac{d C}{d f_i} \cdot \frac{d f_i}{d x_j}`. Both the partial differentiation and the multiplication have to be performed by :meth:`Op.grad`. PyTensor currently imposes the following constraints on the values returned by the :meth:`Op.grad` method: 1) They must be :class:`Variable` instances. 2) When they are types that have dtypes, they must never have an integer dtype. The output gradients passed *to* :meth:`Op.grad` will also obey these constraints. Integers are a tricky subject. Integers are the main reason for having :class:`DisconnectedType`, :class:`NullType` or zero gradient. When you have an integer as an argument to your :meth:`Op.grad` method, recall the definition of a derivative to help you decide what value to return: :math:`\frac{d f}{d x} = \lim_{\epsilon \rightarrow 0} (f(x+\epsilon)-f(x))/\epsilon`. Suppose your function f has an integer-valued output. For most functions you're likely to implement in PyTensor, this means your gradient should be zero, because :math:`f(x+epsilon) = f(x)` for almost all :math:`x`. (The only other option is that the gradient could be undefined, if your function is discontinuous everywhere, like the rational indicator function) Suppose your function :math:`f` has an integer-valued input. This is a little trickier, because you need to think about what you mean mathematically when you make a variable integer-valued in PyTensor. Most of the time in machine learning we mean ":math:`f` is a function of a real-valued :math:`x`, but we are only going to pass in integer-values of :math:`x`". In this case, :math:`f(x+\epsilon)` exists, so the gradient through :math:`f` should be the same whether :math:`x` is an integer or a floating point variable. Sometimes what we mean is ":math:`f` is a function of an integer-valued :math:`x`, and :math:`f` is only defined where :math:`x` is an integer." Since :math:`f(x+\epsilon)` doesn't exist, the gradient is undefined. Finally, many times in PyTensor, integer valued inputs don't actually affect the elements of the output, only its shape. If your function :math:`f` has both an integer-valued input and an integer-valued output, then both rules have to be combined: - If :math:`f` is defined at :math:`x + \epsilon`, then the input gradient is defined. Since :math:`f(x+\epsilon)` would be equal to :math:`f(x)` almost everywhere, the gradient should be zero (first rule). - If :math:`f` is only defined where :math:`x` is an integer, then the gradient is undefined, regardless of what the gradient with respect to the output is. Examples: 1) :math:`f(x,y)` is a dot product between :math:`x` and :math:`y`. :math:`x` and :math:`y` are integers. Since the output is also an integer, :math:`f` is a step function. Its gradient is zero almost everywhere, so :meth:`Op.grad` should return zeros in the shape of :math:`x` and :math:`y`. 2) :math:`f(x,y)` is a dot product between :math:`x` and :math:`y`. :math:`x` is floating point and :math:`y` is an integer. In this case the output is floating point. It doesn't matter that :math:`y` is an integer. We consider :math:`f` to still be defined at :math:`f(x,y+\epsilon)`. The gradient is exactly the same as if :math:`y` were floating point. 3) :math:`f(x,y)` is the argmax of :math:`x` along axis :math:`y`. The gradient with respect to :math:`y` is undefined, because :math:`f(x,y)` is not defined for floating point :math:`y`. How could you take an argmax along a fractional axis? The gradient with respect to :math:`x` is 0, because :math:`f(x+\epsilon, y) = f(x)` almost everywhere. 4) :math:`f(x,y)` is a vector with :math:`y` elements, each of which taking on the value :math:`x` The :meth:`Op.grad` method should return :class:`DisconnectedType` for :math:`y`, because the elements of :math:`f` don't depend on :math:`y`. Only the shape of :math:`f` depends on :math:`y`. You probably also want to implement a connection_pattern method to encode this. 5) :math:`f(x) = int(x)` converts float :math:`x` into an integer. :math:`g(y) = float(y)` converts an integer :math:`y` into a float. If the final cost :math:`C = 0.5 * g(y) = 0.5 g(f(x))`, then the gradient with respect to :math:`y` will be 0.5, even if :math:`y` is an integer. However, the gradient with respect to :math:`x` will be 0, because the output of :math:`f` is integer-valued. .. function:: connection_pattern(node): Sometimes needed for proper operation of :func:`pytensor.gradient.grad`. Returns a list of list of booleans. ``Op.connection_pattern[input_idx][output_idx]`` is true if the elements of ``inputs[input_idx]`` have an effect on the elements of ``outputs[output_idx]``. The ``node`` parameter is needed to determine the number of inputs. Some :class:`Op`\s such as :class:`Subtensor` take a variable number of inputs. If no connection_pattern is specified, :func:`pytensor.gradient.grad` will assume that all inputs have some elements connected to some elements of all outputs. This method conveys two pieces of information that are otherwise not part of the PyTensor graph: 1) Which of the :class:`Op`'s inputs are truly ancestors of each of the :class:`Op`'s outputs. Suppose an :class:`Op` has two inputs, :math:`x` and :math:`y`, and outputs :math:`f(x)` and :math:`g(y)`. :math:`y` is not really an ancestor of :math:`f`, but it appears to be so in the PyTensor graph. 2) Whether the actual elements of each input/output are relevant to a computation. For example, the shape :class:`Op` does not read its input's elements, only its shape metadata. :math:`\frac{d shape(x)}{dx}` should thus raise a disconnected input exception (if these exceptions are enabled). As another example, the elements of the :class:`Alloc` :class:`Op`'s outputs are not affected by the shape arguments to the :class:`Alloc` :class:`Op`. Failing to implement this function for an :class:`Op` that needs it can result in two types of incorrect behavior: 1) :func:`pytensor.gradient.grad` erroneously raising a ``TypeError`` reporting that a gradient is undefined. 2) :func:`pytensor.gradient.grad` failing to raise a ``ValueError`` reporting that an input is disconnected. Even if connection_pattern is not implemented correctly, if :func:`pytensor.gradient.grad` returns an expression, that expression will be numerically correct. .. function:: R_op(inputs, eval_points) Optional, to work with :func:`pytensor.gradient.R_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`. ``inputs`` are the symbolic variables corresponding to the value of the input where you want to evaluate the Jacobian, and ``eval_points`` are the symbolic variables corresponding to the value you want to right multiply the Jacobian with. Same conventions as for the :meth:`Op.grad` method hold. If your :class:`Op` is not differentiable, you can return None. Note that in contrast to the method :meth:`Op.grad`, for :meth:`Op.R_op` you need to return the same number of outputs as there are outputs of the :class:`Op`. You can think of it in the following terms. You have all your inputs concatenated into a single vector :math:`x`. You do the same with the evaluation points (which are as many as inputs and of the shame shape) and obtain another vector :math:`v`. For each output, you reshape it into a vector, compute the Jacobian of that vector with respect to :math:`x` and multiply it by :math:`v`. As a last step you reshape each of these vectors you obtained for each outputs (that have the same shape as the outputs) back to their corresponding shapes and return them as the output of the :meth:`Op.R_op` method. :ref:`List of op with r op support `.