Factor out shared functions

Author: jessegrabowski

pytensor/tensor/optimize.py

@@ -128,6 +128,32 @@ def _find_optimization_parameters(objective: TensorVariable, x: TensorVariable):
     ]
 
 
+def _get_parameter_grads_from_vector(
+    grad_wrt_args_vector: Variable,
+    x_star: Variable,
+    args: Sequence[Variable],
+    output_grad: Variable,
+):
+    """
+    Given a single concatenated vector of objective function gradients with respect to raveled optimization parameters,
+    returns the contribution of each parameter to the total loss function, with the unraveled shape of the parameter.
+    """
+    cursor = 0
+    grad_wrt_args = []
+
+    for arg in args:
+        arg_shape = arg.shape
+        arg_size = arg_shape.prod()
+        arg_grad = grad_wrt_args_vector[:, cursor : cursor + arg_size].reshape(
+            (*x_star.shape, *arg_shape)
+        )
+
+        grad_wrt_args.append(dot(output_grad, arg_grad))
+        cursor += arg_size
+
+    return grad_wrt_args
+
+
 class ScipyWrapperOp(Op, HasInnerGraph):
     """Shared logic for scipy optimization ops"""
 
@@ -348,34 +374,25 @@ def L_op(self, inputs, outputs, output_grads):
 
         implicit_f = grad(inner_fx, inner_x)
 
-        df_dx = atleast_2d(concatenate(jacobian(implicit_f, [inner_x]), axis=-1))
+        df_dx, *df_dtheta_columns = jacobian(
+            implicit_f, [inner_x, *inner_args], disconnected_inputs="ignore"
+        )
 
         df_dtheta = concatenate(
-            [
-                atleast_2d(x, left=False)
-                for x in jacobian(implicit_f, inner_args, disconnected_inputs="ignore")
-            ],
+            [atleast_2d(jac_col, left=False) for jac_col in df_dtheta_columns],
             axis=-1,
         )
 
         replace = dict(zip(self.fgraph.inputs, (x_star, *args), strict=True))
 
-        df_dx_star, df_dtheta_star = graph_replace([df_dx, df_dtheta], replace=replace)
+        df_dx_star, df_dtheta_star = graph_replace(
+            [atleast_2d(df_dx), df_dtheta], replace=replace
+        )
 
         grad_wrt_args_vector = solve(-df_dx_star, df_dtheta_star)
-
-        cursor = 0
-        grad_wrt_args = []
-
-        for arg in args:
-            arg_shape = arg.shape
-            arg_size = arg_shape.prod()
-            arg_grad = grad_wrt_args_vector[:, cursor : cursor + arg_size].reshape(
-                (*x_star.shape, *arg_shape)
-            )
-
-            grad_wrt_args.append(dot(output_grad, arg_grad))
-            cursor += arg_size
+        grad_wrt_args = _get_parameter_grads_from_vector(
+            grad_wrt_args_vector, x_star, args, output_grad
+        )
 
         return [zeros_like(x), *grad_wrt_args]
 
@@ -504,19 +521,9 @@ def L_op(
         df_dx_star, df_dtheta_star = graph_replace([df_dx, df_dtheta], replace=replace)
 
         grad_wrt_args_vector = solve(-df_dx_star, df_dtheta_star)
-
-        cursor = 0
-        grad_wrt_args = []
-
-        for arg in args:
-            arg_shape = arg.shape
-            arg_size = arg_shape.prod()
-            arg_grad = grad_wrt_args_vector[:, cursor : cursor + arg_size].reshape(
-                (*x_star.shape, *arg_shape)
-            )
-
-            grad_wrt_args.append(dot(output_grad, arg_grad))
-            cursor += arg_size
+        grad_wrt_args = _get_parameter_grads_from_vector(
+            grad_wrt_args_vector, x_star, args, output_grad
+        )
 
         return [zeros_like(x), *grad_wrt_args]
 
@@ -529,11 +536,7 @@ def root(
 ):
     """Find roots of a system of equations using scipy.optimize.root."""
 
-    args = [
-        arg
-        for arg in truncated_graph_inputs([equations], [variables])
-        if (arg is not variables and not isinstance(arg, Constant))
-    ]
+    args = _find_optimization_parameters(equations, variables)
 
     root_op = RootOp(variables, *args, equations=equations, method=method, jac=jac)