Add regression test for #1586

Author: jessegrabowski

tests/tensor/test_optimize.py

@@ -8,6 +8,7 @@
 from pytensor.graph import Apply, Op, Type
 from pytensor.tensor import alloc, scalar, scalar_from_tensor, tensor_from_scalar
 from pytensor.tensor.optimize import minimize, minimize_scalar, root, root_scalar
+from pytensor.tensor.special import betaln
 from tests import unittest_tools as utt
 
 
@@ -523,3 +524,44 @@ def L_op(self, inputs, outputs, output_gradients):
         _grad_wrt_num_theta = pt.grad(x_star, num_theta, disconnected_inputs="raise")
     # np.testing.assert_allclose(grad_wrt_num_theta.eval({x: np.pi, num_theta: np.e, str_theta: ":)"}), -1)
     # np.testing.assert_allclose(grad_wrt_num_theta.eval({x: np.pi, num_theta: np.e, str_theta: ":("}), 1)
+
+
+def test_vectorize_root_gradients():
+    """Regression test for https://github.com/pymc-devs/pytensor/issues/1586"""
+    n, a, b = pt.scalars("n a b".split())
+    w_min, w_max = pt.scalars("w_min w_max".split())
+
+    w_support = pt.linspace(w_min, w_max, n + 1)
+
+    k = pt.floor(w_support)
+    ln_n_choose_k = -pt.log(n + 1) - betaln(n - k + 1, k + 1)
+    q_probs = pt.exp(ln_n_choose_k + betaln(k + a, n - k + b) - betaln(a, b))
+
+    c = pt.dscalar("c")  # Unemployment benefit
+    β = pt.dscalar("β")  # Discount rate
+
+    # initial value function guess
+    v0 = pt.dvector("v0")
+
+    # Fixed-point operator
+    T = pt.maximum(w_support / (1 - β), c + β * pt.dot(v0, q_probs))
+
+    v_star, _ = pt.optimize.root(equations=T - v0, variables=v0, method="hybr")
+
+    w_bar = (1 - β) * (c + β * pt.dot(v_star, q_probs))
+
+    # We want to study the impact of change in unemployment and patience on the reserve wage
+    w_grads = pt.grad(w_bar, [c, β])
+
+    c_grid = pt.dmatrix("c_grid")
+    β_grid = pt.dmatrix("β_grid")
+
+    w_bar_grid, *w_grad_grid = pytensor.graph.vectorize_graph(
+        [w_bar, *w_grads], {β: β_grid, c: c_grid}
+    )
+
+    _ = pytensor.function(
+        [v0, c_grid, β_grid, n, a, b, w_min, w_max],
+        [w_bar_grid, *w_grad_grid],
+        on_unused_input="ignore",
+    )