Implement vectorized jacobian and allow arbitrary expression shapes

doc/tutorial/gradients.rst

@@ -101,18 +101,20 @@ PyTensor implements the :func:`pytensor.gradient.jacobian` macro that does all
 that is needed to compute the Jacobian. The following text explains how
 to do it manually.
 
+Using Scan
+----------
+
 In order to manually compute the Jacobian of some function ``y`` with
-respect to some parameter ``x`` we need to use `scan`. What we
-do is to loop over the entries in ``y`` and compute the gradient of
+respect to some parameter ``x`` we can use `scan`.
+In this case, we loop over the entries in ``y`` and compute the gradient of
 ``y[i]`` with respect to ``x``.
 
 .. note::
 
     `scan` is a generic op in PyTensor that allows writing in a symbolic
     manner all kinds of recurrent equations. While creating
     symbolic loops (and optimizing them for performance) is a hard task,
-    effort is being done for improving the performance of `scan`. We
-    shall return to :ref:`scan<tutloop>` later in this tutorial.
+    efforts are being made to improving the performance of `scan`.
 
 >>> import pytensor
 >>> import pytensor.tensor as pt
@@ -124,9 +126,9 @@ do is to loop over the entries in ``y`` and compute the gradient of
 array([[ 8.,  0.],
        [ 0.,  8.]])
 
-What we do in this code is to generate a sequence of integers from ``0`` to
-``y.shape[0]`` using `pt.arange`. Then we loop through this sequence, and
-at each step, we compute the gradient of element ``y[i]`` with respect to
+This code generates a sequence of integers from ``0`` to
+``y.shape[0]`` using `pt.arange`. Then it loops through this sequence, and
+at each step, computes the gradient of element ``y[i]`` with respect to
 ``x``. `scan` automatically concatenates all these rows, generating a
 matrix which corresponds to the Jacobian.
 
@@ -139,6 +141,31 @@ matrix which corresponds to the Jacobian.
     ``x`` anymore, while ``y[i]`` still is.
 
 
+Using automatic vectorization
+-----------------------------
+An alternative way to build the Jacobian is to vectorize the graph that computes a single row or colum of the jacobian
+We can use `Lop` or `Rop` (more about it below) to obtain the row or column of the jacobian and `vectorize_graph`
+to vectorize it to the full jacobian matrix.
+
+>>> import pytensor
+>>> import pytensor.tensor as pt
+>>> from pytensor.gradient import Lop
+>>> from pytensor.graph import vectorize_graph
+>>> x = pt.dvector('x')
+>>> y = x ** 2
+>>> row_cotangent = pt.dvector("row_cotangent")  # Helper variable, it will be replaced during vectorization
+>>> J_row = Lop(y, x, row_cotangent)
+>>> J = vectorize_graph(J_row, replace={row_cotangent: pt.eye(x.size)})
+>>> f = pytensor.function([x], J)
+>>> f([4, 4])
+array([[ 8.,  0.],
+       [ 0.,  8.]])
+
+This avoids the overhead of scan, at the cost of higher memory usage if the jacobian expression has large intermediate operations.
+Also, not all graphs are safely vectorizable (e.g., if different rows require intermediate operations of different sizes).
+For these reasons `jacobian` uses scan by default. The behavior can be changed by setting `vectorize=True`.
+
+
 Computing the Hessian
 =====================
 

pytensor/gradient.py

@@ -11,7 +11,7 @@
 import pytensor
 from pytensor.compile.ops import ViewOp
 from pytensor.configdefaults import config
-from pytensor.graph import utils
+from pytensor.graph import utils, vectorize_graph
 from pytensor.graph.basic import Apply, NominalVariable, Variable
 from pytensor.graph.null_type import NullType, null_type
 from pytensor.graph.op import get_test_values
@@ -703,15 +703,15 @@
         grad_dict[var] = g_var
 
     def handle_disconnected(var):
-        message = (
-            "grad method was asked to compute the gradient "
-            "with respect to a variable that is not part of "
-            "the computational graph of the cost, or is used "
-            f"only by a non-differentiable operator: {var}"
-        )
         if disconnected_inputs == "ignore":
-            pass
+            return
         elif disconnected_inputs == "warn":
+            message = (
+                "grad method was asked to compute the gradient "
+                "with respect to a variable that is not part of "
+                "the computational graph of the cost, or is used "
+                f"only by a non-differentiable operator: {var}"
+            )
             warnings.warn(message, stacklevel=2)
         elif disconnected_inputs == "raise":
             message = utils.get_variable_trace_string(var)
@@ -2021,13 +2021,19 @@
 Exception args: {args_msg}"""
 
 
-def jacobian(expression, wrt, consider_constant=None, disconnected_inputs="raise"):
+def jacobian(
+    expression,
+    wrt,
+    consider_constant=None,
+    disconnected_inputs="raise",
+    vectorize=False,
+):
     """
     Compute the full Jacobian, row by row.
 
     Parameters
     ----------
-    expression : Vector (1-dimensional) :class:`~pytensor.graph.basic.Variable`
+    expression :class:`~pytensor.graph.basic.Variable`
         Values that we are differentiating (that we want the Jacobian of)
     wrt : :class:`~pytensor.graph.basic.Variable` or list of Variables
         Term[s] with respect to which we compute the Jacobian
@@ -2051,62 +2057,74 @@
         output, then a zero variable is returned. The return value is
         of same type as `wrt`: a list/tuple or TensorVariable in all cases.
     """
+    from pytensor.tensor.basic import eye
+    from pytensor.tensor.extra_ops import broadcast_to
 
     if not isinstance(expression, Variable):
         raise TypeError("jacobian expects a Variable as `expression`")
 
-    if expression.ndim > 1:
-        raise ValueError(
-            "jacobian expects a 1 dimensional variable as `expression`."
-            " If not use flatten to make it a vector"
-        )
-
     using_list = isinstance(wrt, list)
     using_tuple = isinstance(wrt, tuple)
+    grad_kwargs = {
+        "consider_constant": consider_constant,
+        "disconnected_inputs": disconnected_inputs,
+    }
 
     if isinstance(wrt, list | tuple):
         wrt = list(wrt)
     else:
         wrt = [wrt]
 
     if all(expression.type.broadcastable):
-        # expression is just a scalar, use grad
-        return as_list_or_tuple(
-            using_list,
-            using_tuple,
-            grad(
-                expression.squeeze(),
-                wrt,
-                consider_constant=consider_constant,
-                disconnected_inputs=disconnected_inputs,
-            ),
+        jacobian_matrices = grad(expression.squeeze(), wrt, **grad_kwargs)
+
+    elif vectorize:
+        expression_flat = expression.ravel()
+        row_tangent = _float_ones_like(expression_flat).type("row_tangent")
+        jacobian_single_rows = Lop(expression.ravel(), wrt, row_tangent, **grad_kwargs)
+
+        n_rows = expression_flat.size
+        jacobian_matrices = vectorize_graph(
+            jacobian_single_rows,
+            replace={row_tangent: eye(n_rows, dtype=row_tangent.dtype)},
         )
+        if disconnected_inputs != "raise":
+            # If the input is disconnected from the cost, `vectorize_graph` has no effect on the respective jacobian
+            # We have to broadcast the zeros explicitly here
+            for i, (jacobian_single_row, jacobian_matrix) in enumerate(
+                zip(jacobian_single_rows, jacobian_matrices, strict=True)
+            ):
+                if jacobian_single_row.ndim == jacobian_matrix.ndim:
+                    jacobian_matrices[i] = broadcast_to(
+                        jacobian_matrix, shape=(n_rows, *jacobian_matrix.shape)
+                    )
 
-    def inner_function(*args):
-        idx = args[0]
-        expr = args[1]
-        rvals = []
-        for inp in args[2:]:
-            rval = grad(
-                expr[idx],
-                inp,
-                consider_constant=consider_constant,
-                disconnected_inputs=disconnected_inputs,
+    else:
+
+        def inner_function(*args):
+            idx, expr, *wrt = args
+            return grad(expr[idx], wrt, **grad_kwargs)
+
+        jacobian_matrices, updates = pytensor.scan(
+            inner_function,
+            sequences=pytensor.tensor.arange(expression.size),
+            non_sequences=[expression.ravel(), *wrt],
+            return_list=True,
+        )
+        if updates:
+            raise ValueError(
+                "The scan used to build the jacobian matrices returned a list of updates"
             )
-            rvals.append(rval)
-        return rvals
-
-    # Computing the gradients does not affect the random seeds on any random
-    # generator used n expression (because during computing gradients we are
-    # just backtracking over old values. (rp Jan 2012 - if anyone has a
-    # counter example please show me)
-    jacobs, updates = pytensor.scan(
-        inner_function,
-        sequences=pytensor.tensor.arange(expression.shape[0]),
-        non_sequences=[expression, *wrt],
-    )
-    assert not updates, "Scan has returned a list of updates; this should not happen."
-    return as_list_or_tuple(using_list, using_tuple, jacobs)
+
+    if jacobian_matrices[0].ndim < (expression.ndim + wrt[0].ndim):
+        # There was some raveling or squeezing done prior to getting the jacobians
+        # Reshape into original shapes
+        jacobian_matrices = [
+            jac_matrix.reshape((*expression.shape, *w.shape))
+            for jac_matrix, w in zip(jacobian_matrices, wrt, strict=True)
+        ]
+
+    return as_list_or_tuple(using_list, using_tuple, jacobian_matrices)
 
 
 def hessian(cost, wrt, consider_constant=None, disconnected_inputs="raise"):

pytensor/graph/replace.py

@@ -232,13 +232,13 @@ def vectorize_graph(
 def vectorize_graph(
     outputs: Sequence[Variable],
     replace: Mapping[Variable, Variable],
-) -> Sequence[Variable]: ...
+) -> list[Variable]: ...
 
 
 def vectorize_graph(
     outputs: Variable | Sequence[Variable],
     replace: Mapping[Variable, Variable],
-) -> Variable | Sequence[Variable]:
+) -> Variable | list[Variable]:
     """Vectorize outputs graph given mapping from old variables to expanded counterparts version.
 
     Expanded dimensions must be on the left. Behavior is similar to the functional `numpy.vectorize`.

pytensor/tensor/basic.py

@@ -3081,6 +3081,10 @@ def flatten(x, ndim=1):
     else:
         dims = (-1,)
 
+    if len(dims) == _x.ndim:
+        # Nothing to ravel
+        return _x
+
     x_reshaped = _x.reshape(dims)
     shape_kept_dims = _x.type.shape[: ndim - 1]
     bcast_new_dim = builtins.all(s == 1 for s in _x.type.shape[ndim - 1 :])

tests/tensor/test_basic.py

@@ -3867,35 +3867,22 @@ class TestInferShape(utt.InferShapeTester):
     def test_Flatten(self):
         atens3 = tensor3()
         atens3_val = random(4, 5, 3)
-        for ndim in (3, 2, 1):
+        for ndim in (2, 1):
             self._compile_and_check(
                 [atens3],
                 [flatten(atens3, ndim)],
                 [atens3_val],
                 Reshape,
-                excluding=["local_useless_reshape"],
             )
 
         amat = matrix()
         amat_val = random(4, 5)
-        for ndim in (2, 1):
-            self._compile_and_check(
-                [amat],
-                [flatten(amat, ndim)],
-                [amat_val],
-                Reshape,
-                excluding=["local_useless_reshape"],
-            )
-
-        avec = vector()
-        avec_val = random(4)
         ndim = 1
         self._compile_and_check(
-            [avec],
-            [flatten(avec, ndim)],
-            [avec_val],
+            [amat],
+            [flatten(amat, ndim)],
+            [amat_val],
             Reshape,
-            excluding=["local_useless_reshape"],
         )
 
     def test_Eye(self):