pytensor 2.38.3__tar.gz → 3.0.0__tar.gz

{pytensor-2.38.3/pytensor.egg-info → pytensor-3.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.38.3
+Version: 3.0.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause

{pytensor-2.38.3 → pytensor-3.0.0}/doc/extending/creating_a_numba_jax_op.rst

@@ -484,7 +484,6 @@ Step 4: Write tests
             from pytensor.configdefaults import config
             from tests.link.jax.test_basic import compare_jax_and_py
             from pytensor.graph import FunctionGraph
-            from pytensor.graph.op import get_test_value
 
             def test_jax_CumOp():
                 """Test JAX conversion of the `CumOp` `Op`."""
@@ -492,8 +491,7 @@ Step 4: Write tests
                 # Create a symbolic input for the first input of `CumOp`
                 a = pt.matrix("a")
 
-                # Create test value tag for a
-                a.tag.test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
+                test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
 
                 # Create the output variable
                 out = pt.cumsum(a, axis=0)
@@ -502,12 +500,12 @@ Step 4: Write tests
                 fgraph = FunctionGraph([a], [out])
 
                 # Pass the graph and inputs to the testing function
-                compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
+                compare_jax_and_py(fgraph, [test_value])
 
                 # For the second mode of CumOp
                 out = pt.cumprod(a, axis=1)
                 fgraph = FunctionGraph([a], [out])
-                compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
+                compare_jax_and_py(fgraph, [test_value])
 
         If the variant :class:`CumprodOp` is not implemented, we can add a test for it as follows:
 
@@ -518,12 +516,12 @@ Step 4: Write tests
             def test_jax_CumOp():
                 """Test JAX conversion of the `CumOp` `Op`."""
                 a = pt.matrix("a")
-                a.tag.test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
+                test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
 
                 with pytest.raises(NotImplementedError):
                     out = pt.cumprod(a, axis=1)
                     fgraph = FunctionGraph([a], [out])
-                    compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
+                    compare_jax_and_py(fgraph, [test_value])
 
 
     .. tab-item:: Numba
@@ -543,23 +541,23 @@ Step 4: Write tests
 
         .. code:: python
 
+            import numpy as np
+            from pytensor.configdefaults import config
             from tests.link.numba.test_basic import compare_numba_and_py
             from pytensor.graph import FunctionGraph
-            from pytensor.compile.sharedvalue import SharedVariable
-            from pytensor.graph.basic import Constant
             from pytensor.tensor import extra_ops
+            import pytensor.tensor as pt
+
+            def test_CumOp(axis, mode):
+                val = pt.matrix("val")
+                test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
 
-            def test_CumOp(val, axis, mode):
                 g = extra_ops.CumOp(axis=axis, mode=mode)(val)
                 g_fg = FunctionGraph(outputs=[g])
 
                 compare_numba_and_py(
                     g_fg,
-                    [
-                        i.tag.test_value
-                        for i in g_fg.inputs
-                        if not isinstance(i, SharedVariable | Constant)
-                    ],
+                    [test_value],
                 )
 
 
@@ -601,7 +599,6 @@ Step 4: Write tests
                 # Create a symbolic input for the first input of `CumOp`
                 a = pt.matrix("a", dtype=dtype)
 
-                # Create test value
                 test_value = np.arange(9, dtype=dtype).reshape((3, 3))
 
                 # Create the output variable

{pytensor-2.38.3 → pytensor-3.0.0}/doc/extending/creating_an_op.rst

@@ -74,20 +74,12 @@ possibilities you may encounter or need.  For that refer to
         def perform(self, node: Apply, inputs_storage: list[Any], output_storage: list[list[Any]]) -> None:
             pass
 
-        # Defines the symbolic expression for the L-operator based on the input and output variables
-        # and the output gradient variables. Optional.
-        def L_op(self, inputs: list[Variable], outputs: list[Variable], output_grads: list[Variable]) -> list[Variable]:
+        # Defines the vector-Jacobian product (pullback/reverse-mode AD). Optional.
+        def pullback(self, inputs: list[Variable], outputs: list[Variable], cotangents: list[Variable]) -> list[Variable]:
             pass
 
-        # 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
-
-        # 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]:
+        # Defines the Jacobian-vector product (pushforward/forward-mode AD). Optional.
+        def pushforward(self, inputs: list[Variable], outputs: list[Variable], tangents: list[Variable]) -> list[Variable]:
             pass
 
         # Defines the symbolic expression for the output shape based on the input shapes
@@ -182,36 +174,42 @@ 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.
 
-:meth:`L_op`
-^^^^^^^^^^^^^^^
+:meth:`pullback`
+^^^^^^^^^^^^^^^^^^^
 
-The :meth:`L_op` method is required if you want to differentiate some cost
-whose expression includes your :class:`Op`. The gradient is
-specified symbolically in this method. It takes three arguments ``inputs``, ``outputs`` and
-``output_gradients``, which are both lists of :ref:`variable`\s, and
-those must be operated on using PyTensor's symbolic language. The :meth:`L_op`
-method must return a list containing one :ref:`variable` for each
-input. Each returned :ref:`variable` represents the gradient with respect
-to that input computed based on the symbolic gradients with respect
-to each output.
+The :meth:`pullback` method is required if you want to differentiate some cost
+whose expression includes your :class:`Op`. It implements the vector-Jacobian
+product (VJP) for reverse-mode automatic differentiation.
+
+It takes three arguments ``inputs``, ``outputs`` and ``cotangents``, which are
+all lists of :ref:`variable`\s that must be operated on using PyTensor's
+symbolic language. The :meth:`pullback` method must return a list containing
+one :ref:`variable` for each input. Each returned :ref:`variable` represents
+the cotangent (gradient) 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. Please refer to :meth:`L_op` for a more detailed
-view.
+this method should return a variable of type :class:`DisconnectedType`
+for that input. If the gradient is not implemented or undefined for some
+input, return a variable of type :class:`NullType` for that input
+(see :func:`pytensor.gradient.grad_not_implemented` and
+:func:`pytensor.gradient.grad_undefined`).
 
-:meth:`R_op`
-^^^^^^^^^^^^^^^
-The :meth:`R_op` method is needed if you want :func:`pytensor.gradient.Rop` to
-work with your :class:`Op`.
+:meth:`pushforward`
+^^^^^^^^^^^^^^^^^^^^^^
 
-This function implements the application of the R-operator on the
-function represented by your :class:`Op`. Let's assume that function is :math:`f`,
-with input :math:`x`, applying the R-operator means computing the
-Jacobian of :math:`f` and right-multiplying it by :math:`v`, the evaluation
-point, namely: :math:`\frac{\partial f}{\partial x} v`.
+The :meth:`pushforward` method is needed if you want :func:`pytensor.gradient.pushforward` to
+work with your :class:`Op` with the non-default ``use_op_pushforward=True`` argument.
+It implements the Jacobian-vector product (JVP) for forward-mode automatic
+differentiation.
+
+Given a function :math:`f` with input :math:`x`, the pushforward computes
+:math:`J \dot{x}` where :math:`J = \frac{\partial f}{\partial x}` is the
+Jacobian and :math:`\dot{x}` are the tangent vectors. It takes ``inputs``,
+``outputs``, and ``tangents`` as arguments. Tangent entries of type
+:class:`DisconnectedType` indicate that the corresponding input is not being
+differentiated. ``None`` must not be used to indicate disconnected entries;
+use :class:`DisconnectedType` variables instead.
 
 
 Example: :class:`Op` definition
@@ -253,22 +251,21 @@ Example: :class:`Op` definition
             # The output shape is the same as the input shape
             return input_shapes
 
-        def L_op(self, inputs: list[TensorVariable], outputs: list[TensorVariable], output_grads: list[TensorVariable]):
-            # Symbolic expression for the gradient
+        def pullback(self, inputs: list[TensorVariable], outputs: list[TensorVariable], cotangents: list[TensorVariable]):
+            # Symbolic expression for the vector-Jacobian product
             # For this Op, the inputs and outputs aren't part of the expression
-            # output_grads[0] is a TensorVariable!
-            return [output_grads[0] * 2]
-
-        def R_op(self, inputs: list[TensorVariable], eval_points: list[TensorVariable | None]) -> list[TensorVariable] | None:
-            # R_op can receive None as eval_points.
-            # That means there is no differentiable path through that input
-            # If this imply that you cannot compute some outputs,
-            # return None for those.
-            if eval_points[0] is None:
-                return None
-            # For this Op, the R_op is the same as the L_op
-            outputs = self(inputs)
-            return self.L_op(inputs, outputs, eval_points)
+            # cotangents[0] is a TensorVariable!
+            return [cotangents[0] * 2]
+
+        def pushforward(self, inputs: list[TensorVariable], outputs: list[TensorVariable], tangents: list[TensorVariable]):
+            from pytensor.gradient import DisconnectedType, disconnected_type
+
+            # pushforward receives DisconnectedType for non-differentiable inputs.
+            # If this means you cannot compute some outputs, return disconnected_type() for those.
+            if isinstance(tangents[0].type, DisconnectedType):
+                return [disconnected_type()]
+            # For this Op, the Jacobian-vector product is symmetric
+            return [tangents[0] * 2]
 
     doubleOp1 = DoubleOp1()
 
@@ -486,8 +483,8 @@ Furthermore, this :class:`Op` will work with batched dimensions, meaning we can
 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.
+We need to be careful in the :meth:`pullback` method, as one of output cotangents 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:`pullback` method, so we don't need to worry about that case.
 
 .. testcode::
 
@@ -534,23 +531,23 @@ If both outputs are disconnected PyTensor will not bother calling the :meth:`L_o
             out2_shape = y_shapes[:-1]
             return [out1_shape, out2_shape]
 
-        def L_op(self, inputs, outputs, output_grads):
+        def pullback(self, inputs, outputs, cotangents):
             x, y = inputs
-            out1_grad, out2_grad = output_grads
+            out1_ct, out2_ct = cotangents
 
-            if isinstance(out1_grad.type, DisconnectedType):
-                x_grad = disconnected_type()
+            if isinstance(out1_ct.type, DisconnectedType):
+                x_ct = disconnected_type()
             else:
-                # Transpose the last two dimensions of the output gradient
-                x_grad = pt.swapaxes(out1_grad, -1, -2)
+                # Transpose the last two dimensions of the output cotangent
+                x_ct = pt.swapaxes(out1_ct, -1, -2)
 
-            if isinstance(out2_grad.type, DisconnectedType):
-                y_grad = disconnected_type()
+            if isinstance(out2_ct.type, DisconnectedType):
+                y_ct = disconnected_type()
             else:
-                # Broadcast the output gradient to the same shape as y
-                y_grad = pt.broadcast_to(pt.expand_dims(out2_grad, -1), y.shape)
+                # Broadcast the output cotangent to the same shape as y
+                y_ct = pt.broadcast_to(pt.expand_dims(out2_ct, -1), y.shape)
 
-            return [x_grad, y_grad]
+            return [x_ct, y_ct]
 
 Let's test the `Op` evaluation:
 
@@ -659,8 +656,8 @@ 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`
+  :meth:`pullback`
+  :meth:`pushforward`
 
 
 Basic Tests
@@ -763,25 +760,25 @@ If you want to see it fail, you can implement an incorrect gradient
                 rng = rng,
             )
 
-Testing the Rop
-^^^^^^^^^^^^^^^
+Testing the pushforward
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-The class :class:`RopLopChecker` defines the methods
-:meth:`RopLopChecker.check_mat_rop_lop`, :meth:`RopLopChecker.check_rop_lop` and :meth:`RopLopChecker.check_nondiff_rop`.
-These allow to test the implementation of the :meth:`R_op` method of a particular :class:`Op`.
+The class :class:`PushforwardPullbackChecker` defines the methods
+:meth:`PushforwardPullbackChecker.check_mat_pushforward_pullback`, :meth:`PushforwardPullbackChecker.check_pushforward_pullback` and :meth:`PushforwardPullbackChecker.check_nondiff_pushforward`.
+These allow to test the implementation of the :meth:`pushforward` method of a particular :class:`Op`.
 
-For instance, to verify the :meth:`R_op` method of the ``DoubleOp``, you can use this:
+For instance, to verify the :meth:`pushforward` method of the ``DoubleOp``, you can use this:
 
 .. testcode:: tests
 
    import numpy
    import tests
-   from tests.test_rop import RopLopChecker
+   from tests.test_rop import PushforwardPullbackChecker
 
-   class TestDoubleOpRop(RopLopChecker):
+   class TestDoubleOpPushforward(PushforwardPullbackChecker):
 
-       def test_double_rop(self):
-           self.check_rop_lop(DoubleOp()(self.x), self.in_shape)
+       def test_double_pushforward(self):
+           self.check_pushforward_pullback(DoubleOp()(self.x), self.in_shape)
 
 
 Running Your Tests
@@ -798,7 +795,7 @@ Modify and execute to compute: ``x * y``.
 
 Modify and execute the example to return two outputs: ``x + y`` and `jx - yj`.
 
-You can omit the :meth:`Rop` functions. Try to implement the testing apparatus described above.
+You can omit the :meth:`pushforward` method. Try to implement the testing apparatus described above.
 
 :download:`Solution<extending_pytensor_solution_1.py>`
 
@@ -830,7 +827,7 @@ It takes an optional :meth:`infer_shape` parameter that must have this signature
 
 .. note::
 
-    As no L_op is defined, this means you won't be able to
+    As no pullback is defined, this means you won't be able to
     differentiate paths that include this :class:`Op`.
 
 .. note::

{pytensor-2.38.3 → pytensor-3.0.0}/doc/extending/op.rst

@@ -244,21 +244,28 @@ Optional methods or attributes
 If you want your :class:`Op` to work with :func:`pytensor.gradient.grad` you also
 need to implement the functions described below.
 
-Gradient
-========
+Automatic Differentiation
+=========================
 
-These are the function required to work with :func:`pytensor.gradient.grad`.
+These are the functions required to work with :func:`pytensor.gradient.grad`
+and :func:`pytensor.gradient.pushforward`.
 
-.. function:: grad(inputs, output_gradients)
+.. function:: pullback(inputs, outputs, cotangents)
 
-  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.
+  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`
@@ -269,108 +276,50 @@ These are the function required to work with :func:`pytensor.gradient.grad`.
   :func:`pytensor.gradient.grad_undefined` and
   :func:`pytensor.gradient.grad_not_implemented`, respectively.
 
-  If an element of ``output_gradient`` is of type
+  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.grad` should
+  the computation of only disconnected outputs, then :meth:`Op.pullback` should
   return :class:`DisconnectedType` variables for those inputs.
 
-  If the :meth:`Op.grad` method is not defined, then PyTensor assumes it has been
+  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`.
 
-  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.
+  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 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`.
+  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.grad` method:
+  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 gradients passed *to* :meth:`Op.grad` will also obey these constraints.
+  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.grad` method, recall the definition of
+  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`.
@@ -410,7 +359,7 @@ These are the function required to work with :func:`pytensor.gradient.grad`.
 
   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
+     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
@@ -423,7 +372,7 @@ These are the function required to work with :func:`pytensor.gradient.grad`.
      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
+     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.
@@ -433,6 +382,29 @@ These are the function required to work with :func:`pytensor.gradient.grad`.
      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`.
@@ -477,32 +449,4 @@ These are the function required to work with :func:`pytensor.gradient.grad`.
   :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.