diffapqp 0.1.0__py3-none-any.whl

lao/diffapqp/layer.py

@@ -0,0 +1,702 @@
+"""
+Layer: differentiates the optimizer solution map x*(theta).
+"""
+from __future__ import annotations
+
+import os
+import torch
+import cvxpy as cp
+from . import value_layer as value_layer_module
+from .value_layer import ValueFuncLayer
+from .solve_result import LayerSolveSummary
+from ..utils import flatten_fortran_torch_batch, unflatten_fortran_torch_batch
+import numpy as np
+import scipy.linalg as la
+import scipy.sparse as sp
+from scipy.sparse import bmat, diags
+from scipy.sparse.linalg import lsqr, minres, spsolve
+from typing import Any
+import warnings
+import multiprocessing as mp
+from threadpoolctl import threadpool_limits
+
+
+class _LokyPoolAdapter:
+    """Expose the Pool methods used by Layer over a Loky process executor."""
+
+    def __init__(self, processes: int, initializer, initargs: tuple):
+        from loky import ProcessPoolExecutor
+
+        self._executor = ProcessPoolExecutor(
+            max_workers=processes,
+            initializer=initializer,
+            initargs=initargs,
+        )
+
+    def map(self, function, iterable, chunksize: int = 1):
+        return list(
+            self._executor.map(
+                function,
+                iterable,
+                chunksize=chunksize,
+            )
+        )
+
+    def submit(self, function, *args, **kwargs):
+        return self._executor.submit(function, *args, **kwargs)
+
+    def close(self) -> None:
+        self._executor.shutdown(wait=True)
+
+    def join(self) -> None:
+        # shutdown(wait=True) has already joined all workers.
+        return None
+
+
+def _as_dense_matrix(matrix):
+    return matrix.toarray() if sp.issparse(matrix) else np.asarray(matrix)
+
+
+def _orthonormal_row_basis(matrix, rank_tol: float):
+    """Return an orthonormal basis for the row space of ``matrix``."""
+    matrix = _as_dense_matrix(matrix)
+    if matrix.size == 0 or matrix.shape[0] == 0:
+        return np.zeros((matrix.shape[1], 0))
+
+    q_mat, r_mat = la.qr(matrix.T, mode="economic", pivoting=False)
+    diag = np.abs(np.diag(r_mat))
+    if diag.size == 0:
+        rank = 0
+    else:
+        threshold = rank_tol * max(matrix.shape) * max(1.0, diag[0])
+        rank = int(np.count_nonzero(diag > threshold))
+    return q_mat[:, :rank]
+
+
+def _rank_prune_active_ineq_indices(A_row_basis, G, active_idx, rank_tol: float):
+    """Keep active inequality rows independent relative to equality rows."""
+    if len(active_idx) == 0:
+        return active_idx
+
+    G_active = _as_dense_matrix(G[active_idx, :])
+    if A_row_basis.shape[1] > 0:
+        G_active = G_active - (G_active @ A_row_basis) @ A_row_basis.T
+
+    _, r_mat, pivots = la.qr(G_active.T, mode="economic", pivoting=True)
+    diag = np.abs(np.diag(r_mat))
+    if diag.size == 0:
+        rank = 0
+    else:
+        threshold = rank_tol * max(G_active.shape) * max(1.0, diag[0])
+        rank = int(np.count_nonzero(diag > threshold))
+    return np.sort(active_idx[pivots[:rank]])
+
+
+class Layer(ValueFuncLayer):
+    
+    # Used for resolving active-set tolerances from solver arguments
+    _ABS_TOL_KEYS = ("eps_abs", "tol_gap_abs", "tol_feas", "tol_feas_abs")
+    _REL_TOL_KEYS = ("eps_rel", "tol_gap_rel", "tol_feas_rel")
+    # Internal switch for reduced-KKT active-row rank pruning.
+    _PRUNE_ACTIVE_INEQ = False
+    _ACTIVE_RANK_TOL = 1e-10
+
+    @staticmethod
+    def _build_regularized_problem(problem: cp.Problem, regularization_eps: float) -> cp.Problem:
+        """Return a problem with a small quadratic term added to all variables.
+        This problem will be used as the new problem for the layer."""
+        if regularization_eps <= 0:
+            return problem
+
+        reg_term = 0
+        for var in problem.variables():
+            reg_term += cp.sum_squares(var)
+
+        objective_expr = problem.objective.args[0] + 0.5 * regularization_eps * reg_term
+        return cp.Problem(cp.Minimize(objective_expr), problem.constraints)
+
+    @classmethod
+    def _infer_solver_tolerance(cls, solver_args: dict, candidate_keys: tuple[str, ...]) -> float | None:
+        for key in candidate_keys:
+            value = solver_args.get(key)
+            if value is not None:
+                return float(value)
+        return None
+
+    @classmethod
+    def _resolve_active_set_tolerances(
+        cls,
+        solver_args: dict,
+        act_ineq_tol: float | None,
+        act_ineq_tol_rel: float | None,
+    ) -> tuple[float, float]:
+        """Choose active-set tolerances from explicit values or solver tolerances.
+        If None, then the tolerance is inferred from the solver arguments."""
+        solver_args = {} if solver_args is None else dict(solver_args)
+        multiplier = 10.0
+        fallback_abs = 1e-4 # default if nothing can be inferred
+        fallback_rel = 1e-4
+
+        inferred_abs = cls._infer_solver_tolerance(solver_args, cls._ABS_TOL_KEYS)
+        inferred_rel = cls._infer_solver_tolerance(solver_args, cls._REL_TOL_KEYS)
+
+        resolved_abs = (
+            float(act_ineq_tol)  # if given, use the given value
+            if act_ineq_tol is not None
+            # if not given, use the inferred value or the fallback value
+            else max(fallback_abs, multiplier * inferred_abs) if inferred_abs is not None else fallback_abs 
+        )
+        resolved_rel = (
+            float(act_ineq_tol_rel)
+            if act_ineq_tol_rel is not None
+            else max(fallback_rel, multiplier * inferred_rel) if inferred_rel is not None else fallback_rel
+        )
+        return resolved_abs, resolved_rel
+
+    def __init__(
+        self,
+        problem: cp.Problem,
+        max_n_cores: int = os.cpu_count(),
+        regularization_eps: float = 0.0,
+        act_ineq_tol: float | None = None,
+        act_ineq_tol_rel: float | None = None,
+    ) -> None:
+        """
+        Initialize the Layer.
+        Args:
+            problem: CVXPY problem with named parameters and variables.
+            max_n_cores: Maximum number of cores to use for parallel solving.
+            regularization_eps: Regularization parameter for the problem (useful for LP).
+            act_ineq_tol: Absolute tolerance for active inequality constraints, default is None.
+            act_ineq_tol_rel: Relative tolerance for active inequality constraints, default is None.
+        """
+        
+        if regularization_eps < 0:
+            raise ValueError("regularization_eps must be nonnegative.")
+        if act_ineq_tol is not None and act_ineq_tol < 0:
+            raise ValueError("act_ineq_tol must be nonnegative.")
+        if act_ineq_tol_rel is not None and act_ineq_tol_rel < 0:
+            raise ValueError("act_ineq_tol_rel must be nonnegative.")
+
+        self.regularization_eps = float(regularization_eps)
+        self.act_ineq_tol = None if act_ineq_tol is None else float(act_ineq_tol)
+        self.act_ineq_tol_rel = None if act_ineq_tol_rel is None else float(act_ineq_tol_rel)
+        self.original_problem = problem
+        self.regularized_problem = self._build_regularized_problem(problem, self.regularization_eps)
+        
+        # use the regularized problem to initialize the ValueFuncLayer
+        super().__init__(self.regularized_problem, max_n_cores=max_n_cores)
+        if self.max_n_cores > 1 and self.backward_pool is None:
+            # Loky provides process parallelism without requiring a __main__
+            # entry-point guard on Windows or macOS. Linux never enters this
+            # branch and retains the existing multiprocessing fork pool.
+            try:
+                self.backward_pool = _LokyPoolAdapter(
+                    processes=self.max_n_cores,
+                    initializer=value_layer_module._init_worker_backward,
+                    initargs=(self.prob_data,),
+                )
+                self.backward_pool_start_method = "loky"
+            except BaseException:
+                self.close()
+                raise
+
+        self._A_row_basis = (
+            _orthonormal_row_basis(self.prob_data["A"], self._ACTIVE_RANK_TOL)
+            if self._PRUNE_ACTIVE_INEQ
+            else None
+        )
+
+    
+    def _build_solution_summary(self, results: list, standard_param_dict_numpy: dict,
+                                canon_time: float, solve_time: float,
+                                act_ineq_tol: float, act_ineq_tol_rel: float) -> LayerSolveSummary:
+        
+        # Reuse the ValueFuncLayer solve path, then identify active inequalities
+        # needed by the reduced-KKT adjoint option.
+        (x_batch, dual_eq_batch, dual_ineq_batch,
+         _, internal_soln_batch, time_batch) = self._summarize_results(
+             results, canon_time, solve_time
+         )  # results: List of SampleSolveResult from _solve_prepared_single_sample; 
+            # canon_time: time for canonicalization; solve_time: time for solving the problems
+        
+        # Obtain the active inequality set for each sample
+        # Standard-form inequality constraints: Gx <= h + sum_i H_i theta_i.
+        Gx = self.prob_data["G"] @ x_batch.T # (n_ineq, batch_size)
+        h = self.prob_data["h"][:, None].copy()
+        for H_key, H_value in self.prob_data["H"].items():
+            h = h + H_value @ standard_param_dict_numpy[H_key].T  # RHS: h += H_i y_i (n_ineq, batch_size)
+        
+        batch_size = len(results)
+        # Per-sample active inequality indices. The relative tolerance keeps the
+        # test scale-aware when right-hand sides have different magnitudes.
+        active_ineq_idx_list = [
+            np.where(
+                np.abs(Gx[:, i] - h[:, i])
+                <= np.maximum(act_ineq_tol, act_ineq_tol_rel * (1 + np.abs(h[:, i])))
+            )[0]
+            for i in range(batch_size)
+        ]
+        if self._PRUNE_ACTIVE_INEQ:
+            active_ineq_idx_list = [
+                _rank_prune_active_ineq_indices(
+                    self._A_row_basis,
+                    self.prob_data["G"],
+                    active_idx,
+                    self._ACTIVE_RANK_TOL,
+                )
+                for active_idx in active_ineq_idx_list
+            ]
+        
+        return LayerSolveSummary(
+            primal=x_batch,
+            dual_eq=dual_eq_batch,
+            dual_ineq=dual_ineq_batch,
+            raw_solutions=internal_soln_batch,
+            timing=time_batch,
+            active_ineq_indices=active_ineq_idx_list,
+            inequality_rhs=h,
+        )
+    
+    def forward(self,
+                param_dict: dict,
+                idx_list: list=None,
+                warm_start: bool = False,
+                update: bool = False,
+                solver: str = "OSQP",
+                verbose: bool = False,
+                solver_args: dict = {},
+                reduced_kkt: bool = False, 
+                mode: str = None,
+                mode_args: dict = {}
+                ):
+        """
+        Forward pass of the Layer.
+        Args:
+            param_dict: Dictionary of parameters, keys must match self.prob.param_dict, values are torch.Tensors.
+            idx_list: List of indices of the current batch of samples to solve. Used for self-defined warm start
+            warm_start: Whether to use warm start.
+            update: Whether to update the problem.
+            solver: The solver to use, default is "OSQP".
+            verbose: Whether to print the solver information.
+            solver_args: Additional arguments passed to the solver.
+            reduced_kkt: Whether to use the reduced KKT system, default is False.
+            mode: The method to use to solve the adjoint system, default is None.
+                Choose from "lsqr": least squares method, "spsolve": sparse solver, "minres": minimum residual method
+                If mode is None, the default is "spsolve".
+            mode_args: Dictionary of arguments for the solver for solving adjoint system
+        Returns:
+            Dictionary of primal variables in their original CVXPY shapes,
+            dual variables for the standard-form problem,
+            raw solver-native outputs,
+            timing information
+        """
+        
+        # Stage 1: solve the batched optimization problems in standard form.
+        # This returns:
+        # - original_param_list_tensor: original-shaped torch inputs kept for autograd
+        # - standard_param_dict_numpy: flattened NumPy parameters used by CVXPY
+        # - results/canon_time/solve_time: standardized primal-dual solve outputs
+        # The actual differentiation of the solution map happens later through
+        # the custom autograd wrapper.
+
+        (
+            results,
+            solution_cache_list,
+            batch_size,
+            original_param_list_tensor,
+            standard_param_dict_numpy,  # Flattened
+            canon_time,
+            solve_time,
+        ) = self.forward_solve(
+            param_dict,
+            idx_list,
+            warm_start,
+            update,
+            solver,
+            verbose,
+            solver_args,
+        )
+
+        # Stage 2: summarize the standardized solve outputs and identify the
+        # active inequality set used by the reduced-KKT adjoint option.
+        act_ineq_tol, act_ineq_tol_rel = self._resolve_active_set_tolerances(
+            solver_args,
+            self.act_ineq_tol,
+            self.act_ineq_tol_rel,
+        )
+        
+        summary = self._build_solution_summary(
+            results,
+            standard_param_dict_numpy,
+            canon_time,
+            solve_time,
+            act_ineq_tol,
+            act_ineq_tol_rel,
+        )
+        
+        # Cache solver-native iterates only for solvers that expose the custom
+        # warm-start/update interface in the modified CVXPY fork.
+        if warm_start and solver.upper() in self._modified_solver_list:
+            self._update_solution_cache(solver, batch_size, idx_list, summary.raw_solutions)
+        
+        # Stage 3: choose how to solve the backward adjoint system.
+        if mode is None:
+            if reduced_kkt:
+                mode = "minres"
+            else:
+                mode = "lsqr"
+            # Direct sparse-solve default used in recent experiments:
+            # mode = "spsolve"
+        elif mode == "minres" and not reduced_kkt:
+            raise ValueError("minres mode is only supported for reduced KKT system.")
+        elif mode == "lsqr" and reduced_kkt:
+            warnings.warn("For the reduced KKT system, minres is usually preferred.")
+        
+        # Stage 4: wrap the already-computed primal solution in a custom
+        # autograd Function. The forward pass simply returns the primal
+        # variables in their original CVXPY shapes; the backward pass solves
+        # the adjoint system to recover parameter gradients.
+        autograd_func = _SolutionMapAutogradFn(
+            primal_batch=summary.primal,
+            dual_ineq_batch=summary.dual_ineq,
+            # _compute_adjoint expects one RHS vector per sample.
+            h_batch=summary.inequality_rhs.T,
+            active_ineq_idx_list_per_sample=summary.active_ineq_indices,
+            prob_data=self.prob_data,
+            ori_var_size_dict=self.ori_var_size_dict,
+            ori_var_shape_dict=self.ori_var_shape_dict,
+            ori_param_shape_dict=self.ori_param_shape_dict,
+            param_keys=list(self.ori_param_shape_dict.keys()),
+            standard_var_size_dict=self.standard_var_size_dict,
+            reduced_kkt=reduced_kkt,
+            mode=mode,
+            mode_args=mode_args,
+            pool=self.backward_pool,
+            max_n_cores=self.max_n_cores,
+        )
+        
+        # The primal_list has been recovered into the ORIGINAL CVXPY variable shapes
+        # [Parameter 1 (batch_size, *), Parameter 2 (batch_size, *), ...]
+        # Link each parameter into autograd graph
+        primal_list_per_var_tensor = autograd_func(*original_param_list_tensor)  
+        
+        # Turn decision variables' list into dictionary
+        # PyTorch autograd Functions only support return tuples/tensors, not dictionaries, 
+        # so rebuild the public variable-name mapping expected by the user-facing API.
+        primal_list_per_var_dict = {}
+        for idx, key in enumerate(self.ori_var_size_dict.keys()):
+            primal_list_per_var_dict[key] = primal_list_per_var_tensor[idx]
+        
+        return primal_list_per_var_dict, summary.dual_eq, summary.dual_ineq, summary.raw_solutions, summary.timing
+        
+def _SolutionMapAutogradFn(primal_batch,
+                        dual_ineq_batch,
+                        h_batch,
+                        active_ineq_idx_list_per_sample,
+                        prob_data,
+                        ori_var_size_dict,
+                        ori_var_shape_dict,
+                        ori_param_shape_dict,
+                        param_keys,
+                        standard_var_size_dict,
+                        reduced_kkt: bool,
+                        mode: str,
+                        mode_args: dict,
+                        max_n_cores: int,
+                        pool: Any):
+    
+    """
+    Define the custom autograd function for the solution-map Layer.
+    
+    Args:
+        primal_batch: Batch of primal variables in the standard-form
+        dual_ineq_batch: Batch of dual variables for inequality constraints in the standard-form
+        h_batch: Batch of right-hand side of inequality constraints in the standard-form
+        active_ineq_idx_list_per_sample: List of array of indices of 
+                    the active inequality constraints of each sample 
+        prob_data: Problem data including P, A, G, H, B, Q, h, g
+        ori_var_size_dict: Dictionary of original variable sizes
+        ori_var_shape_dict: Dictionary of original variable shapes
+        ori_param_shape_dict: Dictionary of original parameter shapes
+        param_keys: Parameter names in the same order as the autograd inputs
+        standard_var_size_dict: Dictionary of standard-form variable sizes.
+        mode: The method to use to solve the adjoint system, default is None. 
+            Choose from 
+                - "lsqr": least squares method
+                - "spsolve": sparse solver,
+                - "minres": minimum residual method
+        mode_args: Dictionary of arguments for the solver for solving adjoint system
+        max_n_cores: Maximum number of cores to use
+        pool: Process pool or executor for the backward pass
+    Returns:
+        PyTorch autograd function for computing gradients
+    """
+    
+    class _SolutionMapAutogradFnFn(torch.autograd.Function):
+        
+        """
+        The forward signature seen by PyTorch is:
+            forward pass: list of parameters -> list of original-shaped primal decision variables.
+                include conversion from standard-form variable order to original variable order.
+            backward pass: list of gradients wrt the primal decision variables -> list of gradients wrt the parameters.
+                include conversion from original variable order to standard-form variable order.
+        """
+        
+        @staticmethod
+        def forward(ctx, *param_list):
+            """
+            Forward pass of autograd function
+                - Receive the standard-form solutions from the forward pass (batch_size, flattened size)
+                - Recover the original-shaped primal decision variables [Parameter 1 (batch_size, *), Parameter 2 (batch_size, *), ...]
+                    which is in the same order as the variable order (NOTE: the pytorch autograd does not support dictionary input and output)
+            args:
+                - param_list: original problem parameters. The optimization has
+                    already been solved before this Function is called; these
+                    inputs define the autograd dependency.
+            return:
+                - x_list: list of primal variables of the original problem in the same order as param_list
+            """
+            # store the context for the backward pass
+            ctx.active_ineq_idx_list_per_sample = active_ineq_idx_list_per_sample
+            ctx.prob_data = prob_data
+            ctx.standard_var_size_dict = standard_var_size_dict
+            ctx.ori_var_shape_dict = ori_var_shape_dict
+            ctx.ori_param_shape_dict = ori_param_shape_dict
+            ctx.param_keys = param_keys
+            ctx.reduced_kkt = reduced_kkt
+            ctx.batch_size = len(active_ineq_idx_list_per_sample)
+            ctx.max_n_cores = max_n_cores
+            ctx.primal_batch = primal_batch
+            ctx.dual_ineq_batch = dual_ineq_batch
+            ctx.h_batch = h_batch
+            ctx.mode = mode
+            ctx.mode_args = mode_args
+            ctx.pool = pool
+
+            primal_batch_tensor = torch.from_numpy(primal_batch).to(dtype=param_list[0].dtype)
+            
+            # Recover original CVXPY variable shapes. CVXPY vectorizes matrices
+            # in column-major order, so the shared helper reverses that layout.
+            dim = 0
+            primal_list_per_var_tensor = []
+            output_var_keys = []
+            for key, size in standard_var_size_dict.items():
+                # The standard-form variable order matches the canonicalized data.
+                if key not in ori_var_size_dict:
+                    raise ValueError(
+                        f"Unexpected standard-form variable '{key}'. "
+                        "DiffAPQP requires canonical variables to match the "
+                        "original CVXPY variables exactly."
+                    )
+                x_slice = primal_batch_tensor[:, dim:dim+size]
+                primal_list_per_var_tensor.append(
+                    unflatten_fortran_torch_batch(x_slice, ori_var_shape_dict[key])
+                )
+                output_var_keys.append(key)
+                dim += size
+            
+            ctx.dtype = primal_batch_tensor.dtype
+            ctx.output_var_keys = output_var_keys
+            
+            # Each returned tensor has shape (batch_size, *original_variable_shape).
+            return tuple(primal_list_per_var_tensor) 
+        
+        @staticmethod
+        def backward(ctx, *grad_output):
+            """
+            backward pass
+            `grad_output` has the same shapes as the original decision-variable
+            tensors returned by forward. We flatten them back to the standard
+            CVXPY variable order before solving the adjoint system.
+            """
+            # Retrieve cached solve data and shape metadata.
+            active_ineq_idx_list_per_sample = ctx.active_ineq_idx_list_per_sample
+            prob_data = ctx.prob_data
+            reduced_kkt = ctx.reduced_kkt
+            batch_size = ctx.batch_size
+            max_n_cores = ctx.max_n_cores
+            primal_batch = ctx.primal_batch
+            dual_ineq_batch = ctx.dual_ineq_batch
+            h_batch = ctx.h_batch
+            dtype = ctx.dtype
+            mode = ctx.mode
+            mode_args = ctx.mode_args
+            pool = ctx.pool
+            ori_var_shape_dict = ctx.ori_var_shape_dict
+            ori_param_shape_dict = ctx.ori_param_shape_dict
+            param_keys = ctx.param_keys
+            output_var_keys = ctx.output_var_keys
+            
+            # Flatten shaped decision-variable gradients back to the standard
+            # CVXPY column-major order expected by the adjoint KKT solve.
+            # In a custom torch.autograd.Function, the backward(ctx, *grad_output) arguments follow the exact order of the tensors returned by forward().
+            grad_output_flat = [
+                flatten_fortran_torch_batch(grad, ori_var_shape_dict[key])
+                for grad, key in zip(grad_output, output_var_keys)
+            ] # as defined in the forward pass, the sequence of grad_output is the same as the sequence of output_var_keys
+            grad_output_tensor = torch.cat(grad_output_flat, dim=1) # (batch_size, sum(dim_i))
+            grad_output_numpy = grad_output_tensor.detach().cpu().numpy()
+            
+            # prepare parallel arguments for the adjoint system
+            grad_args = [(grad_output_numpy[i],
+                          primal_batch[i],
+                          dual_ineq_batch[i],         
+                          active_ineq_idx_list_per_sample[i],
+                          h_batch[i],
+                          reduced_kkt,
+                          mode,
+                          mode_args)
+                        for i in range(batch_size)]
+            
+            # Store problem data in this module for serial execution. Parallel
+            # workers get the same data from ValueFuncLayer's pool initializer.
+            global _GLOBAL_PROB_DATA
+            _GLOBAL_PROB_DATA = prob_data
+            
+            if max_n_cores == 1 or pool is None:
+                # list of dicts of gradients wrt the parameters (in standard-form variable order) for each sample
+                grad_params_list = [_compute_adjoint(arg) for arg in grad_args]
+            else:
+                grad_params_list = pool.map(
+                    _compute_adjoint,
+                    grad_args,
+                    # Control how many samples are handed to each worker.
+                    chunksize=max(1, len(grad_args)//(4*max_n_cores))
+                )
+                
+            # Combine gradients across batch dimension
+            combined_grads = {}
+            grad_param_keys = grad_params_list[0].keys()
+            
+            for key in grad_param_keys:
+                # Stack gradients for each parameter across batch dimension
+                combined_grads[key] = torch.from_numpy(
+                    np.stack([grad_params[key] for grad_params in grad_params_list])).to(dtype = dtype)
+            
+            # Return gradients in the same order as input parameters
+            grad_list = [
+                unflatten_fortran_torch_batch(combined_grads[key], ori_param_shape_dict[key])
+                for key in param_keys
+            ]
+            
+            return tuple(grad_list) # this matches the sequence of param_keys
+
+    return _SolutionMapAutogradFnFn.apply
+
+def _compute_adjoint(args):
+    # for preventing BLAS/OpenMP oversubscription in multiprocessing workers
+    with threadpool_limits(limits=1):
+        return _compute_adjoint_impl(args)
+
+def _compute_adjoint_impl(args: tuple) -> dict:
+    """
+    Compute ONE sample's adjoint vector and parameter gradients.
+    
+    To be used in parallel backward pass.
+
+    Args:
+        args: Gradient wrt the primal solution, primal solution, inequality
+            duals, active inequality indices, inequality RHS, reduced-KKT flag,
+            adjoint linear-solver mode, and solver-specific mode kwargs.
+
+    Returns:
+        dict: Gradient of the parameters.
+    """
+    
+    # g is dL/dx in the flattened standard-form variable order.
+    # Everything is in the standard-form
+    g, primal, dual_ineq, active_ineq_idx, h, reduced_kkt, mode, mode_args = args
+    mode_args = {} if mode_args is None else dict(mode_args)
+    if mp.current_process().name != "MainProcess":
+        # Backward process workers are initialized with problem data in
+        # value_layer.py. Linux fork workers may also inherit stale module data
+        # from a previous benchmark run, so always prefer the initializer data.
+        prob_data = value_layer_module._GLOBAL_PROB_DATA
+    else:
+        prob_data = _GLOBAL_PROB_DATA
+    
+    # P, A, and G are sparse standard-form matrices.
+    P = prob_data["P"]
+    A = prob_data["A"]
+    G = prob_data["G"]
+    
+    if not reduced_kkt:
+        # Full complementarity KKT system. This is generally nonsymmetric and
+        # can be ill-conditioned when slacks or duals are close to zero.
+        K = bmat([
+            [P, A.T, G.T],
+            [A, None, None],
+            [diags(dual_ineq) @ G, None, diags(G @ primal - h)]
+        ]).T  # NOTE: the transpose
+        
+    else:
+        # Reduced active-set KKT system. For QPs this is symmetric and avoids
+        # inactive complementarity rows.
+        # Note that it can also be ill-conditioned
+        G_act = G[active_ineq_idx, :]
+        H_act_dict = {key: value[active_ineq_idx, :] for key, value in prob_data["H"].items()}
+        
+        K = bmat([
+                [P, A.T, G_act.T],
+                [A, None, None],
+                [G_act, None, None]
+            ])
+    
+    # Square KKT branches use the common [primal; equality-dual;
+    # inequality-dual] unknown layout, so augment the RHS with zeros for
+    # dual-adjoint components.
+    
+    g_ = np.zeros(K.shape[1])  # set the dual variables to be zero.
+    g_[:len(g)] = g # NOTE: currently differentiation with respect to the dual variables are not supported
+    
+    # Solve the adjoint system
+    if mode == "lsqr":
+        # LSQR can be slow or inaccurate on poorly scaled full KKT systems;
+        # use it as an iterative fallback when direct solves are too costly.
+        lsqr_args = {
+            "atol": 1e-6,
+            "btol": 1e-6,
+            "iter_lim": 20000,
+            "show": False,
+        }
+        lsqr_args.update(mode_args)
+        adjoint = lsqr(K.tocsc(), g_, **lsqr_args)[0]
+    elif mode == "spsolve":
+        spsolve_args = {}
+        spsolve_args.update(mode_args)
+        adjoint = spsolve(K.tocsc(), g_, **spsolve_args)
+    elif mode == "minres":
+        minres_args = {
+            "rtol": 1e-9
+        }
+        minres_args.update(mode_args)
+        adjoint = minres(K.tocsc(), g_, **minres_args)[0]
+    
+    # Convert the adjoint vector into gradients for each canonical parameter.
+    grad_params = {}
+    
+    for key, value in prob_data["H"].items():
+        if not reduced_kkt:
+            # Full KKT data derivative: D_i = [-Q_i; B_i; diag(lambda) H_i].
+            D = bmat(
+                [
+                    [-prob_data['Q'][key]],
+                    [prob_data['B'][key]],
+                    [diags(dual_ineq) @ value]
+                ]
+            )
+        else:
+            # Reduced KKT data derivative: D_i = [-Q_i; B_i; H_active_i].
+            D = bmat(
+                [
+                    [-prob_data['Q'][key]],
+                    [prob_data['B'][key]],
+                    [H_act_dict[key]]
+                ]
+            )
+        grad_params[key] = D.T @ adjoint
+    
+    return grad_params
+        
+