diffapqp 0.1.0__py3-none-any.whl

lao/diffapqp/value_layer.py

@@ -0,0 +1,799 @@
+"""
+Diffapqp.ValueFuncLayer: takes the differentiation of the value function.
+For given APQP,
+- In the forward pass, it solves the continuous problem,
+and outputs the optimal value (value function) and primal/dual variables
+- In the backward pass, it computes the gradients of the VALUE FUNCTION 
+with respect to the parameters for FREE (no implicit differentiation)
+- Everything can be in parallel
+"""
+
+import torch
+import numpy as np
+from ..cvxpy_compat import warn_missing_cvxpy_fork
+from ..utils import (
+    _flattened_size,
+    flatten_fortran_numpy_sample,
+    get_default_solver_args,
+    return_standard_form_in_cvxpy,
+    unflatten_fortran_torch_batch,
+)
+import cvxpy as cp
+from copy import deepcopy
+import os
+import sys
+import warnings
+import time
+from threadpoolctl import threadpool_limits
+from multiprocessing.pool import ThreadPool
+import multiprocessing as mp
+from .solve_result import PreparedSolveData, PreparedSolveTask, SampleSolveResult
+
+
+def _init_worker_backward(prob_data):
+    """Broadcast read-only canonical problem data once per backward worker.
+    Each worker process receives the data instead of sending/pickling the data with every backward call
+    Only used for the Layer. ValueFuncLayer's backward pass is already efficient for serial solving.
+    """
+    global _GLOBAL_PROB_DATA
+    _GLOBAL_PROB_DATA = prob_data
+
+
+def _should_use_fork_backward_pool() -> bool:
+    """Keep the existing fork path where supported, except on macOS."""
+    return sys.platform != "darwin" and "fork" in mp.get_all_start_methods()
+
+
+class ValueFuncLayer(torch.nn.Module):
+    """
+    Differentiable value function layer for continuous APQP problems.
+    """
+
+    _modified_solver_list = {"SCS", "OSQP", "QPALM", "PROXQP"}
+    
+    def __init__(self, problem: cp.Problem, max_n_cores: int = os.cpu_count()) -> None:
+        """
+        Initialize the ValueFuncLayer.
+        
+        Args:
+            problem (cp.Problem): CVXPY problem with named parameters and variables.
+                                Must contain only continuous variables.
+        """
+        super(ValueFuncLayer, self).__init__()
+        
+        # Validate problem has parameters and variables
+        if len(problem.param_dict) < 1:
+            raise ValueError("Problem must have at least one parameter.")
+        if len(problem.var_dict) < 1:
+            raise ValueError("Problem must have at least one variable.")
+            
+        self._canonicalize(problem) # will also check the problem
+        
+        # Forward solves are threaded because the solver calls release most of
+        # the GIL. On Linux, retain the existing fork-based backward pool for
+        # benchmark compatibility. Windows and macOS leave it to SolMapLayer to
+        # create a Loky process pool because ValueFuncLayer does not use it.
+        if max_n_cores == -1:
+            max_n_cores = os.cpu_count()
+        if max_n_cores > 1:
+            self.forward_pool = ThreadPool(processes=max_n_cores)
+            if _should_use_fork_backward_pool():
+                ctx_mp = mp.get_context("fork")
+                self.backward_pool = ctx_mp.Pool(
+                    processes=max_n_cores,
+                    initializer=_init_worker_backward,
+                    initargs=(self.prob_data,)
+                )
+                self.backward_pool_start_method = "fork"
+            else:
+                self.backward_pool = None
+                self.backward_pool_start_method = None
+        else:
+            self.forward_pool = None
+            self.backward_pool = None
+            self.backward_pool_start_method = None
+        
+        self.max_n_cores = max_n_cores
+        
+        # Expensive CVXPY/solver objects are cached by the first batch slot in
+        # self.cvxpy_prob_list; 
+        # cheap warm-start vectors are cached by sample index for all samples.
+        self.solution_cache_dict = {}
+        self.cvxpy_prob_list = None
+
+    def close(self) -> None:
+        """Close persistent worker pools owned by this layer."""
+        if self.backward_pool is not None:
+            self.backward_pool.close()
+            self.backward_pool.join()
+            self.backward_pool = None
+            self.backward_pool_start_method = None
+
+        if self.forward_pool is not None:
+            self.forward_pool.close()
+            self.forward_pool.join()
+            self.forward_pool = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+        return False
+            
+    def _canonicalize(self, problem: cp.Problem) -> None:
+        """
+        Convert the problem to standard APQP form.
+
+        Standard form:
+        min 1/2 x^T P x + (sum Q_i y_i)^T x
+        s.t. Ax == b + sum B_i y_i
+             Gx <= h + sum H_i z_i
+        """
+        # Store original shapes so public outputs/gradients can be reshaped back
+        # after CVXPY canonicalizes everything into flat vectors.
+        self.ori_var_shape_dict = {var.name(): tuple(var.shape) for var in problem.variables()}
+        self.ori_var_size_dict = {var.name(): _flattened_size(var.shape) for var in problem.variables()}
+        self.ori_param_shape_dict = {param.name(): tuple(param.shape) for param in problem.parameters()}
+        self.ori_param_size_dict = {param.name(): _flattened_size(param.shape) for param in problem.parameters()}
+
+        # Convert once to affine-parametric QP standard form:
+        # Q: linear cost, d: objective offset, B: equality RHS, H: inequality RHS.
+        # prob_data includes P, Q, q, c, d, A, G, b, h, B, H
+        self.prob, self.standard_var_size_dict, self.prob_data = return_standard_form_in_cvxpy(problem, sparse=True)
+        self.n_eq_standard = self.prob_data["A"].shape[0]
+        self.n_ineq_standard = self.prob_data["G"].shape[0]
+        
+        if set(self.standard_var_size_dict) != set(self.ori_var_size_dict):
+            raise ValueError(
+                "CVXPY canonicalization introduced or removed variables. "
+                "DiffAPQP requires standard-form variables to match the "
+                "original CVXPY variables exactly. Please reformulate the problem."
+            )
+
+        for key, size in self.ori_var_size_dict.items():
+            if self.standard_var_size_dict[key] != size:
+                raise ValueError(
+                    f"Variable '{key}' changed size during canonicalization. "
+                    "DiffAPQP requires standard-form variable sizes to match "
+                    "the original CVXPY variables. Please reformulate the problem."
+                )
+    
+    def _prepare_forward_data(self, param_dict: dict):
+        """
+        Validate and convert batched parameters for the forward pass.
+
+        Args:
+            param_dict:
+                Mapping from original CVXPY parameter names to batched torch
+                tensors. After the leading batch dimension, each tensor must
+                match the ORIGINAL CVXPY parameter shape:
+                - scalar parameter: ``(batch,)`` or ``(batch, 1)``
+                - vector parameter of shape ``(n,)``: ``(batch, n)``
+                - matrix parameter of shape ``(m, n)``: ``(batch, m, n)``
+
+        Returns:
+            standard_param_dict_numpy:
+                Parameters converted to NumPy and flattened into the
+                standard-form representation expected by the canonicalized
+                CVXPY problem. Matrix parameters are vectorized in Fortran
+                (column-major) order.
+            original_param_list_tensor:
+                Original torch tensors kept in CVXPY parameter order for
+                autograd bookkeeping.
+            batch_size:
+                Shared batch dimension across all parameters.
+        """
+        
+        standard_param_dict_numpy = {} # {parameter_name: numpy_array (batch_size, flattened_size)}
+        original_param_list_tensor = []
+        batch_size = None
+        
+        for key in self.ori_param_size_dict.keys():
+            # Preserve CVXPY's original parameter order; autograd gradients must
+            # be returned in this same order.
+            if key not in param_dict:
+                raise ValueError(f"Missing parameter '{key}' in input param_dict.")
+            if not isinstance(param_dict[key], torch.Tensor):
+                raise ValueError(f"Parameter {key} must be a torch.Tensor")
+            
+            tensor = param_dict[key]
+            original_param_list_tensor.append(tensor)  # record in CVXPY parameter order
+            tensor_np = tensor.detach().cpu().numpy()
+            expected_shape = self.ori_param_shape_dict[key]
+            
+            if tensor_np.ndim == 0:
+                raise ValueError(
+                    f"Parameter '{key}' must include a leading batch dimension. "
+                    f"Expected shape (batch, ...) for original parameter shape {expected_shape}."
+                )
+
+            if batch_size is None:
+                batch_size = tensor_np.shape[0]
+            elif tensor_np.shape[0] != batch_size:
+                raise ValueError(f"Inconsistent batch size for parameter '{key}'.")
+
+            actual_shape = tuple(tensor_np.shape[1:])
+            if len(expected_shape) == 0:
+                # scalar parameter
+                if actual_shape not in [(), (1,)]:
+                    raise ValueError(
+                        f"Parameter '{key}' has invalid shape {tuple(tensor_np.shape)}. "
+                        f"Expected (batch,) or (batch, 1) for scalar parameter."
+                    )
+            elif actual_shape != expected_shape:
+                raise ValueError(
+                    f"Parameter '{key}' has invalid shape {tuple(tensor_np.shape)}. "
+                    f"Expected (batch, {', '.join(str(dim) for dim in expected_shape)}) "
+                    f"to match original parameter shape {expected_shape}."
+                )
+
+            # Standard-form CVXPY parameters are flat vectors. Matrix-shaped
+            # parameters use CVXPY's column-major vectorization convention.
+            if len(expected_shape) == 0:
+                # Supports scalar parameters passed as shape (batch,).
+                tensor_np = tensor_np.reshape(batch_size, 1)
+            elif len(expected_shape) > 1:
+                tensor_np = np.stack(
+                    [flatten_fortran_numpy_sample(tensor_np[i]) for i in range(batch_size)],
+                    axis=0,
+                )
+            standard_param_dict_numpy[key] = tensor_np
+
+        return standard_param_dict_numpy, original_param_list_tensor, batch_size
+    
+    @staticmethod
+    def _ensure_solver_options_in_inverse_data(inverse_data, solver_args: dict) -> None:
+        """Fill missing `solver_options` fields expected by some cvxpy solver inversions.
+        Some CVXPY versions/solver paths may leave inverse_data.solver_options as None, 
+        while the unpack/invert logic expects it to be a dict.
+        """
+        if inverse_data is None:
+            return
+        if isinstance(inverse_data, (list, tuple)):
+            for item in inverse_data:
+                ValueFuncLayer._ensure_solver_options_in_inverse_data(item, solver_args)
+            return
+        if hasattr(inverse_data, "solver_options") and getattr(inverse_data, "solver_options") is None:
+            inverse_data.solver_options = dict(solver_args) if solver_args is not None else {}
+
+    @staticmethod
+    def _prepare_single_sample(
+        cvxpy_prob: cp.Problem,
+        param_dict: dict,
+        warm_start: bool,
+        update: bool,
+        solver: str,
+        warm_start_solution_dict: dict,
+    ) -> PreparedSolveData:
+        """Assign one sample's parameters and create CVXPY solver data.
+        
+        Args:
+            cvxpy_prob: CVXPY problem in standard form returned by return_standard_form_in_cvxpy
+            param_dict: Dictionary of parameters returned by _prepare_forward_data in standard form
+            warm_start: Whether to use warm start
+            update: Whether to update the problem
+            solver: The solver to use
+            warm_start_solution_dict: Dictionary of warm start solution
+        """
+        for key, value in param_dict.items():
+            cvxpy_prob.param_dict[key].value = value
+        warm_start_solution_dict = {} if warm_start_solution_dict is None else warm_start_solution_dict
+        prepare_start_time = time.time()
+        data, chain, inverse_data = cvxpy_prob.get_problem_data(solver)
+        prepare_time = time.time() - prepare_start_time
+        
+        if solver.upper() in ValueFuncLayer._modified_solver_list:
+            # for the modified solvers, we need to set the update and warm_start flags
+            cache_ready = getattr(cvxpy_prob, "_lao_update_cache_ready", False)
+            # The first solve must build the solver workspace. Later solves may
+            # use the custom update path exposed by the modified CVXPY fork.
+            data["update"] = bool(update and cache_ready)
+            effective_warm_start = warm_start and len(warm_start_solution_dict) > 0
+            data["warm_start"] = effective_warm_start
+            if effective_warm_start:
+                data["warm_start_solution_dict"] = warm_start_solution_dict
+
+        return PreparedSolveData(
+            cvxpy_prob=cvxpy_prob,
+            data=data,
+            chain=chain,
+            inverse_data=inverse_data,
+            prepare_time=prepare_time,
+        )
+
+    @staticmethod
+    def _solve_prepared_single_sample(task: PreparedSolveTask) -> SampleSolveResult:
+        """Solve one sample from already-prepared CVXPY data."""
+        cvxpy_prob = task.prepared.cvxpy_prob
+        data = task.prepared.data
+        chain = task.prepared.chain
+        inverse_data = task.prepared.inverse_data
+        prepare_time = task.prepared.prepare_time
+        
+        n_eq_standard = task.n_eq_standard
+        n_ineq_standard = task.n_ineq_standard
+        solver = task.solver
+        warm_start = task.warm_start
+        verbose = task.verbose
+        solver_args = task.solver_args
+
+        worker_start_time = time.time()
+        # This wraps CVXPY's solver interface, including Python/binding overhead
+        # around the solver's own internal timer.
+        solve_via_data_start_time = time.time()
+        soln = chain.solve_via_data(
+            problem=cvxpy_prob,
+            data=data,
+            # For non-modified solvers, warm_start here is the only handle
+            # For modified solvers, warm_start here is ignored
+            warm_start=warm_start,  # The warm-start is passed for the solver that are not modified
+            verbose=verbose,
+            solver_opts=solver_args,
+        )
+        solve_via_data_time = time.time() - solve_via_data_start_time
+
+        ValueFuncLayer._ensure_solver_options_in_inverse_data(inverse_data, solver_args)
+
+        # Unpack primal/dual values back into the standardized CVXPY problem.
+        unpack_start_time = time.time()
+        cvxpy_prob.unpack_results(soln, chain, inverse_data)
+        unpack_time = time.time() - unpack_start_time
+
+        if solver.upper() in ValueFuncLayer._modified_solver_list:
+            # Mark that this batch slot now has a solver workspace eligible for
+            # a future custom update call.
+            cvxpy_prob._lao_update_cache_ready = True
+
+        if cvxpy_prob.status == "optimal_inaccurate":
+            warnings.warn(
+                "The problem is solved in an inaccurate way. Please consider using a more accurate solver."
+            )
+            print(f"The value is: {cvxpy_prob.value}")
+        elif cvxpy_prob.status != "optimal":
+            raise ValueError(f"The problem is not solved. Status: {cvxpy_prob.status}")
+
+        worker_total_time = time.time() - worker_start_time
+
+        # The standard problem has constraints in the order:
+        # Ax == b and Gx <= h; in case no equality and/or inequality constraints are handled.
+        constraint_idx = 0
+        if n_eq_standard > 0:
+            dual_eq = np.asarray(cvxpy_prob.constraints[constraint_idx].dual_value)
+            constraint_idx += 1
+        else:
+            dual_eq = np.zeros((0,), dtype=float)
+
+        if n_ineq_standard > 0:
+            dual_ineq = np.asarray(cvxpy_prob.constraints[constraint_idx].dual_value)
+        else:
+            dual_ineq = np.zeros((0,), dtype=float)
+
+        return SampleSolveResult(
+            primal=cvxpy_prob.variables()[0].value,              # in standard form
+            dual_eq=dual_eq,                                     # in standard form
+            dual_ineq=dual_ineq,                                 # in standard form
+            value=cvxpy_prob.value,                              
+            raw_solution=soln,                                   # for whatever returned by the internal solver
+            timing={
+                "worker_total_time": worker_total_time,
+                "get_problem_data_time": prepare_time,
+                "solve_via_data_time": solve_via_data_time,
+                "unpack_time": unpack_time,
+            },
+        )
+    
+    @staticmethod
+    def _solve_parallel(cvxpy_prob_list: list,    # standardized CVXPY problems, one per batch slot
+                        param_dict_list: list,    # one parameter dictionary per sample
+                        warm_start: bool = True, 
+                        update: bool = False,
+                        solver: str = "OSQP", 
+                        verbose: bool = False,
+                        solver_args: dict = None,
+                        n_eq_standard: int = 0,
+                        n_ineq_standard: int = 0,
+                        warm_start_solution_dict_list: list = None,  # each element is a dictionary with the same keys
+                        pool = None, # thread pool for parallel solving
+                        ) -> tuple:
+        
+        batch_size = len(param_dict_list)
+        solver_args = {} if solver_args is None else solver_args
+
+        warm_list = warm_start_solution_dict_list if warm_start_solution_dict_list is not None else [{} for _ in range(batch_size)]
+        warn_missing_cvxpy_fork(
+            solver,
+            warm_start=warm_start,
+            update=update,
+            warm_start_solution_dict_list=warm_list,
+        )
+        
+        start_time = time.time()
+
+        # Phase 1 (serial): stuff problem data for each sample. This remains
+        # serial because CVXPY problem-data generation mutates problem objects.
+        prepared_list = []
+        for i in range(batch_size):
+            prepared = ValueFuncLayer._prepare_single_sample(
+                cvxpy_prob=cvxpy_prob_list[i],
+                param_dict=param_dict_list[i],
+                warm_start=warm_start,
+                update=update,
+                solver=solver,
+                warm_start_solution_dict=warm_list[i],
+            )
+            prepared_list.append(
+                PreparedSolveTask(
+                    prepared=prepared,
+                    solver=solver,
+                    warm_start=warm_start,
+                    verbose=verbose,
+                    solver_args=solver_args,
+                    n_eq_standard=n_eq_standard,
+                    n_ineq_standard=n_ineq_standard,
+                )
+            )
+        
+        canon_time = time.time() - start_time
+        
+        start_time = time.time()
+        # Phase 2 (parallel): solve using the already-prepared data. Limit
+        # nested BLAS/OpenMP threads so outer sample parallelism does not
+        # oversubscribe the machine.
+        if pool is None:
+            # serial solving
+            results = [ValueFuncLayer._solve_prepared_single_sample(args) for args in prepared_list]
+        else:
+            with threadpool_limits(limits=1):
+                # parallel solving
+                # threadpool_limits is used to limit the nested BLAS/OpenMP threads 
+                # so outer sample parallelism does not oversubscribe the machine.
+                results = pool.map(ValueFuncLayer._solve_prepared_single_sample, prepared_list)
+        
+        solve_time = time.time() - start_time
+        
+        return results, canon_time, solve_time
+    
+    def _summarize_results(self, results: list, canon_time: float, solve_time: float):
+        
+        # Each results is a list of SampleSolveResult from _solve_prepared_single_sample.
+        # For each sample's result, timing includes worker total time, 
+        # get_problem_data_time, solve_via_data_time, unpack_time
+        # of each sample
+        batch_size = len(results)
+        total_x_dim = sum(self.ori_var_size_dict.values())
+        
+        x_batch = np.zeros((batch_size, total_x_dim))
+        dual_eq_batch = np.zeros((batch_size, self.n_eq_standard))
+        dual_ineq_batch = np.zeros((batch_size, self.n_ineq_standard))
+        optimal_value_batch = []
+        internal_soln_batch = []
+        sample_time_batch = []
+        time_batch = {
+            # this is outside the internal solver time
+            "canon_time": canon_time,
+            "solve_time": solve_time,
+        }
+        
+        for i, result in enumerate(results):
+            x_batch[i] = result.primal
+            dual_eq_batch[i] = result.dual_eq
+            dual_ineq_batch[i] = result.dual_ineq
+            optimal_value_batch.append(result.value)
+            internal_soln_batch.append(result.raw_solution)
+            sample_time_batch.append(result.timing)
+
+        time_batch["sample_times"] = sample_time_batch
+        
+        return x_batch, dual_eq_batch, dual_ineq_batch, optimal_value_batch, internal_soln_batch, time_batch
+    
+    def forward_solve(self,
+                param_dict: dict,
+                idx_list: list=None,
+                warm_start: bool = True,
+                update: bool = False,
+                solver: str = "OSQP",
+                verbose: bool = False,
+                solver_args: dict = {}
+                ):
+        
+        """
+        Solve the problem using the targeted solver in its standard form.
+
+        Args:
+            param_dict (dict): Dict of parameters, keys must match self.prob.param_dict, values are torch.Tensors.
+            idx_list (list): List of indices of the current batch of samples to solve.
+            warm_start (bool): Whether to use warm start.
+            update (bool): Whether to update the problem.
+            solver (str): The solver to use, default is "OSQP".
+            verbose (bool): Whether to print the solver information.
+            solver_args (dict): Additional arguments passed to the solver.
+
+        Returns:
+            tuple:
+                results: List of SampleSolveResult objects.
+                solution_cache_list: Warm-start cache entries for the current batch, or None.
+                batch_size: Number of samples in the batch.
+                original_param_list_tensor: Original torch parameters in CVXPY parameter order.
+                standard_param_dict_numpy: Flattened NumPy parameters for the standard-form problem.
+                canon_time: Time spent preparing CVXPY problem data.
+                solve_time: Wall time spent solving prepared samples.
+        """
+        
+        # standard_param_dict_numpy feeds CVXPY; original_param_list_tensor
+        # preserves the original torch inputs and order needed by autograd.
+        standard_param_dict_numpy, original_param_list_tensor, batch_size = self._prepare_forward_data(param_dict)
+        
+        solver = solver.upper()
+        solver_args = {**get_default_solver_args(solver), **solver_args} # call the default and overwrite
+        
+        solution_cache_list = None
+        
+        # Maintain a bounded expensive cache: one standardized CVXPY problem per
+        # batch slot, not one per dataset sample (which is very expensive in memory). 
+        # Warm-start vectors below are cached per sample id and are much cheaper.
+        if self.cvxpy_prob_list is None:
+            self.cvxpy_prob_list = [deepcopy(self.prob) for _ in range(batch_size)]
+
+        cvxpy_prob_list = self.cvxpy_prob_list[:batch_size]
+
+        if warm_start and solver.upper() in self._modified_solver_list:
+            assert idx_list is not None, "idx_list must be provided for warm start."
+            for idx in idx_list:
+                if idx not in self.solution_cache_dict.keys():
+                    # initialize the solution cache for the sample
+                    self.solution_cache_dict[idx] = {}
+
+            # idx_list maps the current batch positions to per-sample warm-start
+            # vectors from previous calls. the sequence must be maintained
+            solution_cache_list = [self.solution_cache_dict[idx] for idx in idx_list]
+
+        # Split batched NumPy parameters into per-sample dictionaries for CVXPY.
+        param_dict_list = [
+            {k: v[i] for k, v in standard_param_dict_numpy.items()}
+            for i in range(batch_size)
+        ]
+        results, canon_time, solve_time = ValueFuncLayer._solve_parallel(
+            cvxpy_prob_list,
+            param_dict_list,
+            warm_start=warm_start,
+            update=update,
+            solver=solver,
+            verbose=verbose,
+            solver_args=solver_args,
+            n_eq_standard=self.n_eq_standard,
+            n_ineq_standard=self.n_ineq_standard,
+            warm_start_solution_dict_list=(
+                solution_cache_list if solution_cache_list is not None else [{} for _ in range(batch_size)]
+            ),
+            pool=self.forward_pool,
+        )
+        
+        return (
+            results,
+            solution_cache_list,
+            batch_size,
+            original_param_list_tensor,
+            standard_param_dict_numpy,
+            canon_time,
+            solve_time,
+        )
+    
+    def _update_solution_cache(self, solver, batch_size, idx_list, internal_soln_batch):
+        for i in range(batch_size):
+            idx = idx_list[i]
+            # Store solver-native iterates for explicit warm starts in the
+            # modified CVXPY solver interfaces.
+            if solver.upper() == "SCS":
+                self.solution_cache_dict[idx] = {
+                    "x": internal_soln_batch[i]["x"],
+                    "y": internal_soln_batch[i]["y"],
+                    "s": internal_soln_batch[i]["s"]
+                }
+            elif solver.upper() == "OSQP":
+                self.solution_cache_dict[idx] = {
+                    "x": internal_soln_batch[i].x,
+                    "y": internal_soln_batch[i].y
+                }
+            elif solver.upper() == "QPALM":
+                self.solution_cache_dict[idx] = {
+                    "x": internal_soln_batch[i].solution.x,
+                    "y": internal_soln_batch[i].solution.y
+                }
+            elif solver.upper() == "PROXQP":
+                self.solution_cache_dict[idx] = {
+                    "x": internal_soln_batch[i].x,
+                    "y": internal_soln_batch[i].y,
+                    "z": internal_soln_batch[i].z
+                }
+            else:
+                raise ValueError(f"Solver {solver} is not supported for warm start.")
+    
+    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 = {}
+                ):
+        
+        (
+            results,
+            solution_cache_list,
+            batch_size,
+            original_param_list_tensor,
+            standard_param_dict_numpy,
+            canon_time,
+            solve_time,
+        ) = self.forward_solve(
+            param_dict,
+            idx_list,
+            warm_start,
+            update,
+            solver,
+            verbose,
+            solver_args,
+        )
+        
+        (x_batch, dual_eq_batch, dual_ineq_batch, optimal_value_batch, 
+            internal_soln_batch, time_batch) = self._summarize_results(results, canon_time, solve_time)
+        if warm_start and solver.upper() in self._modified_solver_list:
+            self._update_solution_cache(solver, batch_size, idx_list, internal_soln_batch)
+        
+        autograd_fn = _ValueFunctionAutogradFn(
+            dual_eq_batch=dual_eq_batch,
+            dual_ineq_batch=dual_ineq_batch,
+            optimal_value_batch=optimal_value_batch,
+            optimal_primal_batch=x_batch,
+            H_dict=self.prob_data["H"],
+            B_dict=self.prob_data["B"],
+            Q_dict=self.prob_data["Q"],
+            d_dict=self.prob_data["d"],
+            ori_param_size_dict=self.ori_param_size_dict,
+            ori_param_shape_dict=self.ori_param_shape_dict,
+        )
+        
+        # Creates the PyTorch autograd connection from the output value to the input parameters:
+        # Custom autograd only supports tuple input and output
+        # No unflattening is needed here
+        optimal_value_batch = autograd_fn(*original_param_list_tensor)
+        
+        # Recover primal variables in their original CVXPY shapes
+        primal_dict = {}
+        dim = 0
+        for key, size in self.ori_var_size_dict.items():
+            shape = self.ori_var_shape_dict[key]
+            x_slice = x_batch[:, dim:dim+size]
+            primal_dict[key] = unflatten_fortran_torch_batch(torch.tensor(x_slice), shape)
+            dim += size
+        if dim != x_batch.shape[1]:
+            raise ValueError(
+                "Standard-form primal dimension does not match the original "
+                "CVXPY variable dimensions. Please reformulate the problem."
+            )
+        
+        dual_eq_batch = torch.tensor(dual_eq_batch)
+        dual_ineq_batch = torch.tensor(dual_ineq_batch)
+        
+        # note that the primal and dual is not differentiable in ValueFuncLayer.forward
+        return optimal_value_batch, primal_dict, dual_eq_batch, dual_ineq_batch, internal_soln_batch, time_batch
+
+def _ValueFunctionAutogradFn(dual_eq_batch, dual_ineq_batch, optimal_value_batch, optimal_primal_batch,
+                             H_dict, B_dict, Q_dict, d_dict, ori_param_size_dict, ori_param_shape_dict):
+    """Define the value-function autograd wrapper.
+    
+    Args:
+        dual_eq_batch: Batch of dual variables for equality constraints of the canonical form problem
+        dual_ineq_batch: Batch of dual variables for inequality constraints of the canonical form problem
+        optimal_value_batch: Batch of optimal objective values
+        optimal_primal_batch: Batch of optimal primal variables in standard form
+        H_dict: Dictionary mapping parameter names to inequality constraint matrices
+        B_dict: Dictionary mapping parameter names to equality constraint matrices
+        Q_dict: Dictionary mapping parameter names to first order objective vector
+        d_dict: Dictionary mapping parameter names to scalar objective-offset vectors
+        ori_param_size_dict: Dictionary of original parameter sizes (flattened)
+    
+    Returns:
+        PyTorch autograd function for computing gradients
+    """
+    class _ValueFunctionAutogradFnFn(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, *param_list):
+            """Return precomputed optimal values while storing data for backward.
+            
+            Args:
+                param_list: List of parameter tensors (dummy inputs since forward pass 
+                        already computed the optimal value and primal/dual)
+                
+            Returns:
+                Tensor of optimal objective values (batch_size,)
+            """
+            # Save the optimal primal/dual solution and parameter-affine data.
+            ctx.H_dict = H_dict
+            ctx.B_dict = B_dict
+            ctx.Q_dict = Q_dict
+            ctx.d_dict = d_dict
+            ctx.dual_ineq_batch = dual_ineq_batch
+            ctx.dual_eq_batch = dual_eq_batch
+            ctx.ori_param_size_dict = ori_param_size_dict
+            ctx.ori_param_shape_dict = ori_param_shape_dict
+            
+            return torch.tensor(optimal_value_batch)
+        
+        @staticmethod
+        def backward(ctx, *grad_output):
+            """Compute value-function gradients by the envelope theorem.
+            
+            Since this custom autograd function has one output, ``grad_output`` is a
+            single-element tuple. ``grad_output[0]`` has the same shape as
+            ``optimal_value_batch``: ``(batch_size,)``.
+
+            Each returned parameter gradient has the same shape as that parameter's input
+            to the forward pass. And the number of the gradients is the same as the number of the parameters.
+            The sequence of the gradients is the same as the sequence of the parameters.
+            
+            However, dictionary-type input and output are not allowed in this custom autograd function.
+            
+            Args:
+                grad_output: Gradient of loss (after this layer) with respect to outputs
+                
+            Returns:
+                Tuple of gradients for each parameter
+            """
+
+            dual_ineq_batch = ctx.dual_ineq_batch  # (batch_size, n_ineq)
+            dual_eq_batch = ctx.dual_eq_batch      # (batch_size, n_eq)
+            H_dict = ctx.H_dict
+            B_dict = ctx.B_dict
+            Q_dict = ctx.Q_dict
+            d_dict = ctx.d_dict
+            ori_param_size_dict = ctx.ori_param_size_dict
+            ori_param_shape_dict = ctx.ori_param_shape_dict
+            
+            # For standard form, dV/dtheta_i =
+            #   -lambda^T H_i - nu^T B_i + x*^T Q_i + d_i.
+            # Scale by upstream grad_output for general scalar losses of V.
+            grad_list = []
+            for param_name, param_size in ori_param_size_dict.items():
+                grad_value = (
+                    -dual_ineq_batch @ H_dict[param_name]
+                    - dual_eq_batch @ B_dict[param_name]
+                    + optimal_primal_batch @ Q_dict[param_name]
+                    + d_dict[param_name]
+                )
+                grad_flat = torch.einsum(
+                    "b, bj->bj",
+                    grad_output[0], 
+                    torch.as_tensor(
+                        grad_value,
+                        dtype=grad_output[0].dtype,
+                        device=grad_output[0].device,
+                    ),
+                )
+                
+                # recover the gradient to the original parameter shapes
+                param_shape = ori_param_shape_dict[param_name]
+                if len(param_shape) == 0:
+                    grad_list.append(grad_flat[:, 0])
+                elif len(param_shape) == 1:
+                    grad_list.append(grad_flat)
+                else:
+                    # matrix parameter
+                    grad_np = grad_flat.detach().cpu().numpy()
+                    grad_reshaped = np.stack(
+                        [grad_np[i].reshape(param_shape, order="F") for i in range(grad_np.shape[0])],
+                        axis=0,
+                    )
+                    grad_list.append(
+                        torch.as_tensor(
+                            grad_reshaped,
+                            dtype=grad_output[0].dtype,
+                            device=grad_output[0].device,
+                        )
+                    )
+            
+            return tuple(grad_list)
+
+    return _ValueFunctionAutogradFnFn.apply