
.. _op_contract:

=============
:class:`Op`\s
=============

An :class:`Op` is a :ref:`graph object <graphstructures>` 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 <views_and_inplace>` 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.

Automatic Differentiation
=========================

These are the functions required to work with :func:`pytensor.gradient.grad`
and :func:`pytensor.gradient.pushforward`.

.. function:: pullback(inputs, outputs, cotangents)

  Implements the vector-Jacobian product (VJP) for reverse-mode automatic
  differentiation. This is the primary method for gradient computation.

  Given a function :math:`f` with inputs :math:`x` and outputs :math:`y = f(x)`,
  the pullback computes :math:`\bar{x} = \bar{y} J` where
  :math:`J = \frac{\partial f}{\partial x}` is the Jacobian and
  :math:`\bar{y}` are the cotangent (row) vectors (upstream gradients).

  Both ``inputs``, ``outputs``, and ``cotangents`` are lists of symbolic PyTensor
  :class:`Variable`\s and those must be operated on using PyTensor's symbolic
  language. The :meth:`Op.pullback` method must return a list containing one
  :class:`Variable` for each input. Each returned :class:`Variable` represents
  the cotangent with respect to that input computed based on the symbolic
  cotangents 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 ``cotangents`` 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.pullback` should
  return :class:`DisconnectedType` variables for those inputs.

  If :meth:`Op.pullback` is not defined, then PyTensor assumes it has been
  forgotten.  Symbolic differentiation will fail on a graph that
  includes this :class:`Op`.

  If an :class:`Op` has a single vector-valued output ``y`` and a single
  vector-valued input ``x``, then :meth:`Op.pullback` will be passed ``x``,
  ``y``, and a cotangent vector ``z``. Define ``J`` to be the Jacobian of ``y``
  with respect to ``x``. The method should return ``dot(z, J)``. When
  :func:`pytensor.grad` calls :meth:`Op.pullback`, 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(z, J)`` is the gradient of C with
  respect to ``x``. If there are other :class:`Op`\s that act on ``x``,
  :func:`pytensor.grad` will add up the terms of ``x``'s gradient contributed
  by each :meth:`Op.pullback`.

  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 vector-Jacobian product.
  Note that an :class:`Op`'s inputs and outputs may include scalars,
  matrices, sparse arrays, or higher-dimensional tensors — not just vectors.
  The VJP contract still applies: conceptually, each of these can be viewed
  as a vector that has been reshaped into a tensor. The returned cotangent
  for each input must have the same shape as that input, and each incoming
  cotangent in ``cotangents`` will have the same shape as the corresponding
  output.

  PyTensor currently imposes the following constraints on the values
  returned by the :meth:`Op.pullback` 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 cotangents passed *to* :meth:`Op.pullback` 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.pullback` 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.pullback` 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.pullback` 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:: pushforward(inputs, outputs, tangents)

  Implements the Jacobian-vector product (JVP) for forward-mode automatic
  differentiation.

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

  ``inputs`` are the symbolic variables corresponding to the input values,
  ``outputs`` are the symbolic variables corresponding to the output values,
  and ``tangents`` are the symbolic variables corresponding to the tangent
  vectors to right-multiply the Jacobian with. Tangent entries of type
  :class:`DisconnectedType` indicate that the corresponding input is not being
  differentiated.

  The method must return the same number of outputs as there are outputs of
  the :class:`Op`. For each output, the result is the Jacobian of that output
  with respect to the inputs, right-multiplied by the tangent vector. Return
  a variable of type :class:`DisconnectedType` for outputs that are disconnected
  from all inputs.

.. 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.


