Add helper to build hessian product

Co-authored-by: Adrian Seyboldt <[email protected]>

Author: ricardoV94

pytensor/gradient.py

@@ -2052,6 +2052,72 @@ def hessian(cost, wrt, consider_constant=None, disconnected_inputs="raise"):
     return as_list_or_tuple(using_list, using_tuple, hessians)
 
 
+def hessian_prod(cost, wrt, p, **grad_kwargs):
+    """Return the expression of the Hessian times a vector p.
+
+    Parameters
+    ----------
+    cost: Scalar (0-dimensional) variable.
+    wrt: Vector (1-dimensional tensor) 'Variable' or list of Vectors
+    p: Vector (1-dimensional tensor) 'Variable' or list of Vectors
+        Each vector will be used for the hessp wirt to exach input variable
+    **grad_kwargs:
+        Keyword arguments passed to `grad` function.
+
+    Returns
+    -------
+    :class:` Vector or list of Vectors
+        The Hessian times p of the `cost` with respect to (elements of) `wrt`.
+
+    Examples
+    --------
+
+    .. testcode::
+
+        import numpy as np
+        from scipy.optimize import minimize
+
+        from pytensor import function
+        from pytensor.tensor import vector
+        from pytensor.gradient import jacobian, hessian_prod
+
+        x = vector('x')
+        p = vector('p')
+
+        rosen = (100 * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2).sum()
+        rosen_hessp = hessian_prod(rosen, x, p)
+        rosen_jac = jacobian(rosen, x)
+
+        rosen_fn = function([x], rosen)
+        rosen_jac_fn = function([x], rosen_jac)
+        rosen_hessp_fn = function([x, p], rosen_hessp)
+        x0 = np.array([1.3, 0.7, 0.8, 1.9, 1.2])
+        res = minimize(
+            rosen_fn,
+            x0,
+            method="Newton-CG",
+            jac=rosen_jac_fn,
+            hessp=rosen_hessp_fn,
+            options={"xatol": 1e-8},
+        )
+        assert res.success
+        np.testing.assert_allclose(res.x, np.ones_like(x0), atol=1e-3)
+
+
+    """
+    wrt_list = wrt if isinstance(wrt, Sequence) else [wrt]
+    p_list = p if isinstance(p, Sequence) else [p]
+    grad_wrt_list = grad(cost, wrt=wrt_list, **grad_kwargs)
+    Hp_list = [
+        grad(grad_wrt @ p, wrt=wrt, **grad_kwargs)
+        for grad_wrt, p in zip(grad_wrt_list, p_list, strict=True)
+    ]
+
+    if isinstance(wrt, Variable):
+        return Hp_list[0]
+    return Hp_list
+
+
 def _is_zero(x):
     """
     Returns 'yes', 'no', or 'maybe' indicating whether x

tests/test_gradient.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+from scipy.optimize import rosen_hess_prod
 
 import pytensor
 import pytensor.tensor.basic as ptb
@@ -22,6 +23,7 @@
     grad_scale,
     grad_undefined,
     hessian,
+    hessian_prod,
     jacobian,
     subgraph_grad,
     zero_grad,
@@ -1081,3 +1083,49 @@ def test_jacobian_disconnected_inputs():
     func_s = pytensor.function([s2], jacobian_s)
     val = np.array(1.0).astype(pytensor.config.floatX)
     assert np.allclose(func_s(val), np.zeros(1))
+
+
+def test_hessp():
+    x = vector("x", dtype="float64")
+    rosen = (100 * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2).sum()
+    p = vector("p")
+
+    rosen_hess_prod_pt = hessian_prod(rosen, wrt=x, p=p)
+
+    x_test = 0.1 * np.arange(9)
+    p_test = 0.5 * np.arange(9)
+    np.testing.assert_allclose(
+        rosen_hess_prod_pt.eval({x: x_test, p: p_test}),
+        rosen_hess_prod(x_test, p_test),
+    )
+
+
+def test_hessp_example():
+    import numpy as np
+    from scipy.optimize import minimize
+
+    from pytensor import function
+    from pytensor.gradient import hessian_prod, jacobian
+    from pytensor.tensor import vector
+
+    x = vector("x")
+    p = vector("p")
+
+    rosen = (100 * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2).sum()
+    rosen_hessp = hessian_prod(rosen, x, p)
+    rosen_jac = jacobian(rosen, x)
+
+    rosen_fn = function([x], rosen)
+    rosen_jac_fn = function([x], rosen_jac)
+    rosen_hessp_fn = function([x, p], rosen_hessp)
+    x0 = np.array([1.3, 0.7, 0.8, 1.9, 1.2])
+    res = minimize(
+        rosen_fn,
+        x0,
+        method="Newton-CG",
+        jac=rosen_jac_fn,
+        hessp=rosen_hessp_fn,
+        options={"xatol": 1e-8},
+    )
+    assert res.success
+    np.testing.assert_allclose(res.x, np.ones_like(x0), atol=1e-3)