Tweak ScipyVectorWrapperOp compute_implicit_gradients for readability

pytensor/tensor/optimize.py

@@ -361,7 +361,7 @@ class ScipyVectorWrapperOp(ScipyWrapperOp):
 
         Notes
         -----
-        The gradents are computed using the implicit function theorem. Given a fuction `f(x, theta) = 0`, and a function
+        The gradients are computed using the implicit function theorem. Given a function `f(x, theta) = 0`, and a function
         `x_star(theta)` that, given input parameters theta returns `x` such that `f(x_star(theta), theta) = 0`, we can
         compute the gradients of `x_star` with respect to `theta` as follows:
 
@@ -387,8 +387,12 @@ class ScipyVectorWrapperOp(ScipyWrapperOp):
         fgraph = self.fgraph
         inner_x, *inner_args = self.inner_inputs
         implicit_f = self.inner_outputs[0]
+        if is_minimization:
+            # The implicit function in minimization is grad(x, theta) == 0
+            implicit_f = grad(implicit_f, inner_x)
 
-        df_dx, *arg_grads = grad(
+        # Call grad to see what arguments are connected
+        _, *arg_grads = grad(
             implicit_f.sum(),
             [inner_x, *inner_args],
             disconnected_inputs="ignore",
@@ -404,9 +408,6 @@ class ScipyVectorWrapperOp(ScipyWrapperOp):
             # No differentiable arguments, return disconnected/null gradients
             return arg_grads
 
-        if is_minimization:
-            implicit_f = grad(implicit_f, inner_x)
-
         # Gradients are computed using the inner graph of the optimization op, not the actual inputs/outputs of the op.
         packed_inner_args, packed_arg_shapes, implicit_f = pack_inputs_of_objective(
             implicit_f,
@@ -426,11 +427,11 @@ class ScipyVectorWrapperOp(ScipyWrapperOp):
 
         # Replace inner inputs (abstract dummies) with outer inputs (the actual user-provided symbols)
         # at the solution point. Innner arguments aren't needed anymore, delete them to avoid accidental references.
-        del inner_x
-        del inner_args
+        del inner_x, inner_args
         inner_to_outer_map = tuple(zip(fgraph.inputs, (x_star, *args)))
         df_dx_star, df_dtheta_star = graph_replace(
-            [df_dx, df_dtheta], inner_to_outer_map
+            [df_dx, df_dtheta],
+            replace=inner_to_outer_map,
         )
 
         if df_dtheta_star.ndim == 0 or df_dx_star.ndim == 0:
@@ -455,13 +456,11 @@ class ScipyVectorWrapperOp(ScipyWrapperOp):
         for i, (arg, to_diff) in enumerate(zip(args, args_to_diff)):
             if not to_diff:
                 # Store the null grad we got from the initial `grad` call
-                null_grad = arg_grads[i]
-                assert isinstance(null_grad.type, NullType | DisconnectedType)
-                final_grads.append(null_grad)
-                continue
-
+                g = arg_grads[i]
+                assert isinstance(g.type, NullType | DisconnectedType)
+            else:
+                # Compute non-null grad and chain with output_grad
                 arg_grad = next(grad_wrt_args_iter)
-
                 if arg_grad.ndim > 0 and output_grad.ndim > 0:
                     g = tensordot(output_grad, arg_grad, [[0], [0]])
                 else: