source_current/src/ffolayer/ffocp_eq.py

import os
import copy
from concurrent.futures import ThreadPoolExecutor

from contextlib import contextmanager
from threadpoolctl import threadpool_limits

@contextmanager
def _limit_native_threads(n: int = 1):
    if threadpool_limits is None:
        yield
    else:
        with threadpool_limits(limits=n):
            yield

import cvxpy as cp
import numpy as np
import torch

try:
    from cvxtorch import TorchExpression
except Exception:
    TorchExpression = None

def _require_cvxtorch():
    if TorchExpression is None:
        raise ImportError(
            "cvxtorch is required for this feature. Install it with:\n"
            "  pip install git+https://github.com/cvxpy/cvxtorch.git"
        )

from cvxpy.constraints.exponential import ExpCone
from cvxpy.constraints.psd import PSD
from cvxpy.constraints.second_order import SOC

from .utils import to_numpy, to_torch, slice_params_for_batch


@torch.no_grad()
def _compare_grads(params_req, grads, ground_truth_grads):
    est_chunks, gt_chunks = [], []
    for p, ge, gg in zip(params_req, grads, ground_truth_grads):
        ge = torch.zeros_like(p) if ge is None else ge.detach()
        gg = torch.zeros_like(p) if gg is None else gg.detach()
        est_chunks.append(ge.reshape(-1))
        gt_chunks.append(gg.reshape(-1))
    est = torch.cat(est_chunks)
    gt = torch.cat(gt_chunks)
    eps = 1e-12
    denom = (est.norm() * gt.norm()).clamp_min(eps)
    cos_sim = torch.dot(est, gt) / denom
    l2_diff = (est - gt).norm()
    return cos_sim, l2_diff


def _cvx_sum_or_zero(terms):
    return cp.sum(terms) if len(terms) > 0 else cp.Constant(0.0)


def _has_pnorm_atom(expr) -> bool:
    try:
        nm_fn = getattr(expr, "name", None)
        if callable(nm_fn):
            nm = nm_fn()
            if nm in {"pnorm", "norm1", "norm_inf"}:
                return True
    except Exception:
        pass

    try:
        cls = expr.__class__.__name__.lower()
        if cls in {"pnorm", "norm1", "norminf", "norm_inf"}:
            return True
    except Exception:
        pass

    for a in getattr(expr, "args", []) or []:
        if _has_pnorm_atom(a):
            return True
    return False


def _infer_objective_expr(problem: cp.Problem):
    obj = problem.objective
    if isinstance(obj, cp.Minimize):
        return obj.expr
    if isinstance(obj, cp.Maximize):
        return -obj.expr
    expr = getattr(obj, "expr", None)
    if expr is None:
        raise ValueError("Unsupported objective type; expected Minimize/Maximize.")
    return expr


def _expcone_dual_dot(u_triplet, c: ExpCone):
    ux, uy, uz = u_triplet
    x, y, z = c.args
    return cp.sum(cp.multiply(ux, x)) + cp.sum(cp.multiply(uy, y)) + cp.sum(cp.multiply(uz, z))


def _split_expcone_dual_value(dv, shapes3):
    if isinstance(dv, (list, tuple)) and len(dv) == 3:
        out = [np.asarray(d, dtype=float) for d in dv]
        for k in range(3):
            if tuple(out[k].shape) != tuple(shapes3[k]):
                if out[k].size == int(np.prod(shapes3[k])):
                    out[k] = out[k].reshape(shapes3[k])
                else:
                    raise ValueError(f"ExpCone dual block {k} shape mismatch: got {out[k].shape}, expected {shapes3[k]}")
        return out

    dv_arr = np.asarray(dv, dtype=float)

    if dv_arr.ndim >= 1 and dv_arr.shape[-1] == 3:
        base = dv_arr.shape[:-1]
        if tuple(base) == tuple(shapes3[0]) and tuple(base) == tuple(shapes3[1]) and tuple(base) == tuple(shapes3[2]):
            return [dv_arr[..., k].reshape(shapes3[k]) for k in range(3)]

    block = int(np.prod(shapes3[0]))
    if int(dv_arr.size) == 3 * block and int(np.prod(shapes3[1])) == block and int(np.prod(shapes3[2])) == block:
        tmp = dv_arr.reshape((block, 3))
        return [tmp[:, k].reshape(shapes3[k]) for k in range(3)]

    raise ValueError(
        f"Cannot parse ExpCone dual_value with shape {dv_arr.shape} into 3 blocks of shapes {shapes3}."
    )


def _active_counts_one(b, ctx, i: int, tol: float):
    out = {}

    out["eq"] = sum(int(np.prod(f.shape)) for f in b["eq_functions"])

    out["ineq"] = sum(
        int(np.sum(np.asarray(ctx.scalar_ineq_slack[j][i]) <= tol))
        for j in range(len(b["scalar_ineq_functions"]))
    )

    soc_cnt = 0
    for c in b["soc_constraints"]:
        t_val = c.args[0].expr.value
        x_val = c.args[1].expr.value
        if t_val is None or x_val is None:
            continue

        t = np.asarray(t_val, dtype=float).reshape(-1)   # (k,) or (1,)
        x = np.asarray(x_val, dtype=float)

        if t.size == 1:
            soc_cnt += int((t.item() - np.linalg.norm(x.ravel())) <= tol)
        else:
            if x.ndim == 1:
                norms = np.full(t.size, np.linalg.norm(x.ravel()))
            elif x.shape[-1] == t.size:
                norms = np.linalg.norm(x.reshape(-1, t.size), axis=0)
            elif x.shape[0] == t.size:
                norms = np.linalg.norm(x.reshape(t.size, -1), axis=1)
            else:
                flat = x.ravel()
                if flat.size % t.size == 0:
                    norms = np.linalg.norm(flat.reshape(t.size, -1), axis=1)
                else:
                    norms = np.full(t.size, np.linalg.norm(flat))

            soc_cnt += int(np.sum((t - norms) <= tol))
    out["soc"] = soc_cnt

    exp_cnt = 0
    for c in b["exp_cones"]:
        x_val, y_val, z_val = (c.args[0].value, c.args[1].value, c.args[2].value)
        if x_val is None or y_val is None or z_val is None:
            continue

        xv = np.asarray(x_val, dtype=float).reshape(-1)  # (k,) or (1,)
        yv = np.asarray(y_val, dtype=float).reshape(-1)
        zv = np.asarray(z_val, dtype=float).reshape(-1)

        k = max(xv.size, yv.size, zv.size)

        # broadcast scalars to vector length k if needed
        if xv.size == 1 and k > 1: xv = np.full(k, xv.item())
        if yv.size == 1 and k > 1: yv = np.full(k, yv.item())
        if zv.size == 1 and k > 1: zv = np.full(k, zv.item())

        with np.errstate(over="ignore", divide="ignore", invalid="ignore"):
            active = np.zeros(k, dtype=bool)
            active[yv <= tol] = True
            mask = yv > tol
            resid = zv[mask] - yv[mask] * np.exp(xv[mask] / yv[mask])
            active[mask] = resid <= tol

        exp_cnt += int(np.sum(active))
    out["exp"] = exp_cnt

    psd_cnt = 0
    for c in b["psd_cones"]:
        X = np.asarray(c.expr.value, dtype=float)
        X = 0.5 * (X + X.T)
        psd_cnt += int(np.linalg.eigvalsh(X).min() <= tol)
    out["psd"] = psd_cnt

    out["cone_total"] = out["soc"] + out["exp"] + out["psd"]
    out["total"] = out["eq"] + out["ineq"] + out["cone_total"]
    return out

def active_counts_dict(ctx, tol = None, reduce: str = "sum"):
    tol = float(ctx.mt.slack_tol if tol is None else tol)
    per_batch = [_active_counts_one(ctx.bundles[i], ctx, i, tol) for i in range(ctx.batch_size)]
    if reduce is None:
        return per_batch
    keys = per_batch[0].keys()
    return {k: sum(d[k] for d in per_batch) for k in keys}


def _build_problem_bundle(
    problem: cp.Problem,
    parameters,
    variables,
    alpha: float,
    dual_cutoff: float,
    slack_tol: float,
    eps: float,
):
    """
    Build and return a dict containing EVERYTHING needed for one problem:
      - forward problem, perturbed problem
      - cvxpy Parameters for dvars/duals/masks
      - torch callables for phi and each dual term
      - pnorm tangent caches + TorchExpression for g (for gradient wrt variables)
    """
    objective_expr = _infer_objective_expr(problem)

    # ---- split constraints ----
    eq_funcs = []
    scalar_ineq_funcs = []
    soc_constraints = []
    exp_cones = []
    psd_cones = []

    for c in problem.constraints:
        if isinstance(c, cp.constraints.zero.Equality):
            eq_funcs.append(c.expr)
        elif isinstance(c, cp.constraints.nonpos.Inequality):
            scalar_ineq_funcs.append(c.expr)
        elif isinstance(c, SOC):
            soc_constraints.append(c)
        elif isinstance(c, ExpCone):
            exp_cones.append(c)
        elif isinstance(c, PSD):
            psd_cones.append(c)
        else:
            raise ValueError(f"Unsupported constraint type: {type(c)}")

    param_order = list(parameters)
    variables = list(variables)

    # ---- original problem (forward) ----
    eq_constraints = [f == 0 for f in eq_funcs]
    scalar_ineq_constraints = [g <= 0 for g in scalar_ineq_funcs]

    forward_problem = cp.Problem(
        cp.Minimize(objective_expr),
        eq_constraints + scalar_ineq_constraints + soc_constraints + exp_cones + psd_cones,
    )

    # ---- dvar params ----
    dvar_params = [cp.Parameter(shape=v.shape) for v in variables]

    # ---- dual params (old) for eq/scalar ineq ----
    eq_dual_params = [cp.Parameter(shape=f.shape) for f in eq_funcs]
    scalar_ineq_dual_params = [cp.Parameter(shape=g.shape, nonneg=True) for g in scalar_ineq_funcs]

    # ---- scalar active masks ----
    scalar_active_mask_params = [cp.Parameter(shape=g.shape, nonneg=True) for g in scalar_ineq_funcs]

    # ---- SOC dual placeholders (old) and linear constraints ----
    soc_dual_params_0 = [cp.Parameter(shape=c.dual_variables[0].shape, nonneg=True) for c in soc_constraints]
    soc_dual_params_1 = [cp.Parameter(shape=c.dual_variables[1].shape) for c in soc_constraints]

    soc_dual_product = _cvx_sum_or_zero([
        cp.multiply(cp.pnorm(c.args[1].expr, p=2) - c.args[0].expr, u)
        for u, c in zip(soc_dual_params_0, soc_constraints)
    ])
    soc_lin_constraints = [
        (soc_dual_params_1[j].T @ soc_constraints[j].args[1].expr
         + cp.multiply(soc_constraints[j].args[0].expr, soc_dual_params_0[j])) == 0
        for j in range(len(soc_constraints))
    ]

    # ---- ExpCone dual placeholders (old) ----
    exp_dual_params = [[cp.Parameter(shape=dv.shape) for dv in c.dual_variables] for c in exp_cones]
    exp_dual_product = _cvx_sum_or_zero([
        _expcone_dual_dot(u3, c) for u3, c in zip(exp_dual_params, exp_cones)
    ])

    # ---- PSD dual placeholders (old) ----
    psd_dual_params = [cp.Parameter(shape=c.dual_variables[0].shape) for c in psd_cones]
    psd_dual_product = _cvx_sum_or_zero([
        cp.sum(cp.multiply(u, c.expr)) for u, c in zip(psd_dual_params, psd_cones)
    ])

    # ---- pnorm tangent support for scalar inequalities (scalar-only) ----
    pnorm_ineq_ids = []
    non_pnorm_scalar_ids = []
    pnorm_xstar_params = []
    pnorm_grad_params = []
    pnorm_tangent_constraints = []
    pnorm_g_torch = []

    for j, g in enumerate(scalar_ineq_funcs):
        is_scalar = int(np.prod(g.shape)) == 1
        is_pnorm = is_scalar and _has_pnorm_atom(g)
        if not is_pnorm:
            non_pnorm_scalar_ids.append(j)
            continue

        local_id = len(pnorm_ineq_ids)
        pnorm_ineq_ids.append(j)

        xs = []
        gs = []
        for v in variables:
            xs.append(cp.Parameter(shape=v.shape))
            gs.append(cp.Parameter(shape=v.shape))
        pnorm_xstar_params.append(xs)
        pnorm_grad_params.append(gs)

        lin = cp.Constant(0.0)
        for v_id, v in enumerate(variables):
            dv = v - pnorm_xstar_params[local_id][v_id]
            lin += cp.sum(cp.multiply(pnorm_grad_params[local_id][v_id], dv))

        pnorm_tangent_constraints.append(cp.multiply(scalar_active_mask_params[j], lin) == 0)

        pnorm_g_torch.append(
            TorchExpression(
                g,
                provided_vars_list=[*variables, *param_order],
            ).torch_expression
        )

    # ---- perturbed problem ----
    vars_dvars_product = _cvx_sum_or_zero([cp.sum(cp.multiply(dv, v)) for dv, v in zip(dvar_params, variables)])
    scalar_ineq_dual_product = _cvx_sum_or_zero([
        cp.sum(cp.multiply(lm, g)) for lm, g in zip(scalar_ineq_dual_params, scalar_ineq_funcs)
    ])

    new_objective = (1.0 / float(alpha)) * vars_dvars_product + objective_expr
    new_objective += scalar_ineq_dual_product + soc_dual_product + exp_dual_product
    # Note: psd_dual_product is intentionally NOT added to the perturbed objective.
    # For PSD cones, keeping them as explicit constraints in the perturbed problem
    # lets the solver handle the conic geometry directly.

    active_eq_constraints = [
        cp.multiply(scalar_active_mask_params[j], scalar_ineq_funcs[j]) == 0
        for j in non_pnorm_scalar_ids
    ]

    perturbed_problem = cp.Problem(
        cp.Minimize(new_objective),
        eq_constraints + active_eq_constraints + soc_lin_constraints + pnorm_tangent_constraints + psd_cones,
    )

    # ---- TorchExpressions for loss pieces (phi and dual terms) ----
    phi_torch = TorchExpression(
        objective_expr,
        provided_vars_list=[*variables, *param_order],
    ).torch_expression

    eq_terms = [cp.sum(cp.multiply(du, f)) for du, f in zip(eq_dual_params, eq_funcs)]
    eq_dual_term_torch = TorchExpression(
        _cvx_sum_or_zero(eq_terms),
        provided_vars_list=[*variables, *param_order, *eq_dual_params],
    ).torch_expression

    ineq_terms = [cp.sum(cp.multiply(du, g)) for du, g in zip(scalar_ineq_dual_params, scalar_ineq_funcs)]
    ineq_dual_term_torch = TorchExpression(
        _cvx_sum_or_zero(ineq_terms),
        provided_vars_list=[*variables, *param_order, *scalar_ineq_dual_params],
    ).torch_expression

    if len(exp_cones) > 0:
        exp_terms = [_expcone_dual_dot(du3, c) for du3, c in zip(exp_dual_params, exp_cones)]
        exp_dual_term_torch = TorchExpression(
            _cvx_sum_or_zero(exp_terms),
            provided_vars_list=[*variables, *param_order, *[u for tri in exp_dual_params for u in tri]],
        ).torch_expression
    else:
        exp_dual_term_torch = None

    if len(psd_cones) > 0:
        psd_terms = [cp.sum(cp.multiply(du, c.expr)) for du, c in zip(psd_dual_params, psd_cones)]
        psd_dual_term_torch = TorchExpression(
            _cvx_sum_or_zero(psd_terms),
            provided_vars_list=[*variables, *param_order, *psd_dual_params],
        ).torch_expression
    else:
        psd_dual_term_torch = None

    non_pnorm_set = set(non_pnorm_scalar_ids)
    pnorm_set = set(pnorm_ineq_ids)
    pnorm_map = {j: lid for lid, j in enumerate(pnorm_ineq_ids)}

    scalar_is_scalar = [int(np.prod(g.shape)) == 1 for g in scalar_ineq_funcs]
    scalar_scalar_indices = [j for j, f in enumerate(scalar_is_scalar) if f]
    scalar_nonscalar_indices = [j for j, f in enumerate(scalar_is_scalar) if not f]

    return dict(
        alpha=float(alpha),
        dual_cutoff=float(dual_cutoff),
        slack_tol=float(slack_tol),
        eps=float(eps),

        param_order=param_order,
        variables=variables,

        objective=objective_expr,
        eq_functions=eq_funcs,
        scalar_ineq_functions=scalar_ineq_funcs,
        scalar_is_scalar=scalar_is_scalar,
        scalar_scalar_indices=scalar_scalar_indices,
        scalar_nonscalar_indices=scalar_nonscalar_indices,
        soc_constraints=soc_constraints,
        exp_cones=exp_cones,
        psd_cones=psd_cones,

        eq_constraints=eq_constraints,
        scalar_ineq_constraints=scalar_ineq_constraints,
        soc_lin_constraints=soc_lin_constraints,
        active_eq_constraints=active_eq_constraints,

        # problems
        problem=forward_problem,
        perturbed_problem=perturbed_problem,

        # cvx params
        dvar_params=dvar_params,
        eq_dual_params=eq_dual_params,
        scalar_ineq_dual_params=scalar_ineq_dual_params,
        scalar_active_mask_params=scalar_active_mask_params,
        soc_dual_params_0=soc_dual_params_0,
        soc_dual_params_1=soc_dual_params_1,
        exp_dual_params=exp_dual_params,
        psd_dual_params=psd_dual_params,

        # pnorm tangent
        pnorm_ineq_ids=pnorm_ineq_ids,
        non_pnorm_scalar_ids=non_pnorm_scalar_ids,
        pnorm_xstar_params=pnorm_xstar_params,
        pnorm_grad_params=pnorm_grad_params,
        pnorm_tangent_constraints=pnorm_tangent_constraints,
        pnorm_g_torch=pnorm_g_torch,

        # precomputed sets/maps
        non_pnorm_set=non_pnorm_set,
        pnorm_set=pnorm_set,
        pnorm_map=pnorm_map,

        # torch callables
        phi_torch=phi_torch,
        eq_dual_term_torch=eq_dual_term_torch,
        ineq_dual_term_torch=ineq_dual_term_torch,
        exp_dual_term_torch=exp_dual_term_torch,
        psd_dual_term_torch=psd_dual_term_torch,
    )


def FFOLayer(
    problem,
    parameters,
    variables,
    alpha: float = 100.0,
    dual_cutoff: float = 1e-3,
    slack_tol: float = 1e-8,
    eps: float = 1e-13,
    compute_cos_sim: bool = False,
    max_workers: int = 8,
    backward_eps: float = 1e-3,
    verbose: bool = False,
):
    _require_cvxtorch()
    print(f"FFOLayer forward eps = {eps}, backward eps = {backward_eps}")
    return _FFOLayer(
        problem=problem,
        parameters=parameters,
        variables=variables,
        alpha=alpha,
        dual_cutoff=dual_cutoff,
        slack_tol=slack_tol,
        eps=eps,
        backward_eps=backward_eps,
        compute_cos_sim=compute_cos_sim,
        max_workers=max_workers,
        verbose=verbose,
    )


class _FFOLayer(torch.nn.Module):
    def __init__(
        self,
        problem,
        parameters,
        variables,
        alpha,
        dual_cutoff,
        slack_tol,
        eps,
        backward_eps,
        compute_cos_sim,
        max_workers: int = 8,
        verbose: bool = False,
    ):
        super().__init__()

        self.alpha = float(alpha)
        self.dual_cutoff = float(dual_cutoff)
        self.slack_tol = float(slack_tol)
        self.eps = float(eps)
        self.backward_eps = float(backward_eps)
        self._compute_cos_sim = bool(compute_cos_sim)
        self.verbose = bool(verbose)

        self._problem_proto = problem

        # If problem is a list, user may pass parameters/variables as list-of-list (one list per problem).
        self._params_list_proto = None
        self._vars_list_proto = None

        if isinstance(problem, (list, tuple)):
            problem_list = list(problem)
            if len(problem_list) == 0:
                raise ValueError("Empty problem_list.")

            # Case A: parameters/variables are list-of-list aligned with problem_list
            if (
                isinstance(parameters, (list, tuple)) and len(parameters) == len(problem_list)
                and len(parameters) > 0 and isinstance(parameters[0], (list, tuple))
            ):
                if not (isinstance(variables, (list, tuple)) and len(variables) == len(problem_list)
                        and len(variables) > 0 and isinstance(variables[0], (list, tuple))):
                    raise ValueError("When problem is a list and parameters is list-of-list, variables must be list-of-list too.")

                self._params_list_proto = [list(pi) for pi in parameters]
                self._vars_list_proto = [list(vi) for vi in variables]

                self._param_templates = list(self._params_list_proto[0])
                self._var_templates = list(self._vars_list_proto[0])

            # Case B: parameters/variables are flat templates; we'll map by name in _lazy_init_from_B
            else:
                self._param_templates = list(parameters)
                self._var_templates = list(variables)
        else:
            self._param_templates = list(parameters)
            self._var_templates = list(variables)
        # self._problem_proto = problem
        # self._param_templates = list(parameters)
        # self._var_templates = list(variables)
        self._max_workers_user = max_workers
        self._initialized = False

        self.num_problems = 0
        self.bundles = None
        self.problem_list = None
        self.perturbed_problem_list = None
        self._ref_param_order = None
        self._ref_vars = None
        self._ws_primal_fwd = None
        self._executor = None

        self.forward_solve_time = 0.0
        self.backward_solve_time = 0.0
        self.forward_setup_time = 0.0
        self.backward_setup_time = 0.0

        self._solver_args_fwd = None
        self._solver_args_bwd = None
    
    def _infer_B_from_params(self, params):
        ref_param_order = self._param_templates
        batch_sizes = []
        for i, (p, qtmpl) in enumerate(zip(params, ref_param_order)):
            if p.ndimension() == qtmpl.ndim:
                bs = 0
            elif p.ndimension() == qtmpl.ndim + 1:
                bs = int(p.size(0))
                if bs <= 0:
                    raise ValueError(f"Parameter {i} has empty batch dimension.")
            else:
                raise ValueError(
                    f"Invalid dim for parameter {i}: got {p.ndimension()}, expected {qtmpl.ndim} or {qtmpl.ndim+1}."
                )

            p_shape = p.shape if bs == 0 else p.shape[1:]
            if tuple(p_shape) != tuple(qtmpl.shape):
                raise ValueError(f"Parameter {i} shape mismatch: expected {qtmpl.shape}, got {p.shape}.")

            batch_sizes.append(bs)

        batch_sizes = np.array(batch_sizes, dtype=int)
        if np.any(batch_sizes > 0):
            nonzero = batch_sizes[batch_sizes > 0]
            B = int(nonzero[0])
            if np.any(nonzero != B):
                raise ValueError(f"Inconsistent batch sizes: {batch_sizes}.")
        else:
            B = 1
        return B


    def _lazy_init_from_B(self, B: int, solver_args: dict):
        if self._initialized:
            return

        if isinstance(self._problem_proto, (list, tuple)):
            problem_list = list(self._problem_proto)
            if len(problem_list) != B:
                raise ValueError(f"Got batch size B={B}, but problem_list has len={len(problem_list)}.")

            # parameters_list = list(self._param_templates)
            # variables_list = list(self._var_templates)
            # if not (len(parameters_list) == len(variables_list) == len(problem_list)):
            #     raise ValueError("When passing problem as list, parameters and variables must be list-of-list aligned.")
            if self._params_list_proto is not None:
                parameters_list = self._params_list_proto
                variables_list = self._vars_list_proto
                if not (len(parameters_list) == len(variables_list) == len(problem_list)):
                    raise ValueError("When passing problem as list, parameters and variables must be list-of-list aligned.")
                # sanity: each inner list length matches template length
                P = len(self._param_templates)
                V = len(self._var_templates)
                for i in range(B):
                    if len(parameters_list[i]) != P:
                        raise ValueError(f"parameters_list[{i}] length mismatch: expected {P}, got {len(parameters_list[i])}")
                    if len(variables_list[i]) != V:
                        raise ValueError(f"variables_list[{i}] length mismatch: expected {V}, got {len(variables_list[i])}")
            else:
                # Otherwise, map by name from each problem
                pnames = [p.name() for p in self._param_templates]
                vnames = [v.name() for v in self._var_templates]
                parameters_list, variables_list = [], []
                for prob_i in problem_list:
                    pmap = prob_i.param_dict
                    vmap = prob_i.var_dict
                    parameters_list.append([pmap[n] for n in pnames])
                    variables_list.append([vmap[n] for n in vnames])
        else:
            pnames = [p.name() for p in self._param_templates]
            vnames = [v.name() for v in self._var_templates]

            problem_list, parameters_list, variables_list = [], [], []
            # Clear solver cache before deepcopy – solver objects are not picklable
            saved_cache = getattr(self._problem_proto, '_solver_cache', None)
            if saved_cache is not None:
                self._problem_proto._solver_cache = {}
            for _ in range(int(B)):
                prob_i = copy.deepcopy(self._problem_proto)
                pmap = prob_i.param_dict
                vmap = prob_i.var_dict
                params_i = [pmap[n] for n in pnames]
                vars_i = [vmap[n] for n in vnames]
                problem_list.append(prob_i)
                parameters_list.append(params_i)
                variables_list.append(vars_i)
            if saved_cache is not None:
                self._problem_proto._solver_cache = saved_cache

        self.num_problems = len(problem_list)
        if self.num_problems == 0:
            raise ValueError("Empty problem_list.")

        self.max_workers = int(self._max_workers_user or min(os.cpu_count() or 1, self.num_problems))
        print(f"max_workers: {self.max_workers}")

        bundles = []
        for prob_i, params_i, vars_i in zip(problem_list, parameters_list, variables_list):
            bundles.append(_build_problem_bundle(
                prob_i,
                parameters=params_i,
                variables=vars_i,
                alpha=self.alpha,
                dual_cutoff=self.dual_cutoff,
                slack_tol=self.slack_tol,
                eps=self.eps,
            ))
        self.bundles = bundles

        self.problem_list = [b["problem"] for b in bundles]
        self.perturbed_problem_list = [b["perturbed_problem"] for b in bundles]

        self._ref_param_order = bundles[0]["param_order"]
        self._ref_vars = bundles[0]["variables"]

        self._ws_cache_fwd = {}  # key -> {scs_x, scs_y, scs_s}
        self._ws_cache_bwd = {}  # key -> {scs_x, scs_y, scs_s}
        self._scs_solvers = {}  # i -> SCS solver instance (for direct SCS path)
        self._scs_data_hash = None  # hash of (A, b) params to detect changes
        self._scs_mapping = None  # {primal_slice, eq_dual_slice, ineq_dual_slice, c_p_slice}

        self._executor = ThreadPoolExecutor(max_workers=self.max_workers)

        self._FFOLayerFn = _make_ffo_fn(self, solver_args=solver_args)
        self._initialized = True


    def close(self):
        ex = getattr(self, "_executor", None)
        if ex is not None:
            ex.shutdown(wait=True)
            self._executor = None

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass


    def forward(self, *params, solver_args=None):
        if solver_args is None:
            solver_args = {}
        solver = solver_args.get("solver", cp.SCS)
        if solver == cp.SCS:
            default_solver_args = dict(
                solver=cp.SCS,
                warm_start=False,
                ignore_dpp=True,
                max_iters=2500,
                eps=self.eps,
                verbose=False,
            )
        else:
            default_solver_args = dict(ignore_dpp=False)

        solver_args = {**default_solver_args, **solver_args}
        self._warm_start = bool(solver_args.get("warm_start", False))
        self._ws_keys = solver_args.pop("ws_keys", None)

        if not self._initialized:
            B = self._infer_B_from_params(params)
            self._lazy_init_from_B(B, solver_args)

        self._solver_args_fwd = dict(solver_args)

        self._solver_args_bwd = dict(solver_args)
        self._solver_args_bwd["max_iters"] = 2500
        self._solver_args_bwd["warm_start"] = True
        if "eps" in self._solver_args_bwd:
            self._solver_args_bwd["eps"] = float(self.backward_eps)

        # Fn = _make_ffo_fn(self, solver_args)
        # return Fn.apply(*params)
        return self._FFOLayerFn.apply(*params)

def _make_ffo_fn(mt: "_FFOLayer", solver_args: dict):
    solver_args = dict(solver_args)
    class _FFOLayerFn(torch.autograd.Function):
        @staticmethod
        def forward(ctx, *params):
            ctx.mt = mt
            ctx.bundles = mt.bundles
            ctx.solver_args = solver_args
            ctx.dtype = params[0].dtype
            ctx.device = params[0].device if isinstance(params[0], torch.Tensor) else 'cpu'

            ref_param_order = mt._ref_param_order
            batch_sizes = []
            for i, (p, qtmpl) in enumerate(zip(params, ref_param_order)):
                if p.dtype != ctx.dtype or p.device != ctx.device:
                    raise ValueError(f"Parameter {i} dtype/device mismatch.")
                if p.ndimension() == qtmpl.ndim:
                    bs = 0
                elif p.ndimension() == qtmpl.ndim + 1:
                    bs = int(p.size(0))
                    if bs <= 0:
                        raise ValueError(f"Parameter {i} has empty batch dimension.")
                else:
                    raise ValueError(f"Invalid dim for parameter {i}: got {p.ndimension()}, expected {qtmpl.ndim} or {qtmpl.ndim+1}.")
                batch_sizes.append(bs)

                p_shape = p.shape if bs == 0 else p.shape[1:]
                if tuple(p_shape) != tuple(qtmpl.shape):
                    raise ValueError(f"Parameter {i} shape mismatch: expected {qtmpl.shape}, got {p.shape}.")

            ctx.batch_sizes = np.array(batch_sizes, dtype=int)
            ctx.batch = bool(np.any(ctx.batch_sizes > 0))
            if ctx.batch:
                nonzero = ctx.batch_sizes[ctx.batch_sizes > 0]
                B = int(nonzero[0])
                if np.any(nonzero != B):
                    raise ValueError(f"Inconsistent batch sizes: {ctx.batch_sizes}.")
            else:
                B = 1
            if ctx.batch and B != mt.num_problems:
                raise ValueError(f"Batch size ({B}) must equal number of problems ({mt.num_problems}).")
            ctx.batch_size = B

            params_np_all = [to_numpy(p) for p in params]

            def _slice_params_np(i: int):
                if ctx.batch:
                    return [arr[i] if bs > 0 else arr for arr, bs in zip(params_np_all, ctx.batch_sizes)]
                return params_np_all

            ref_bundle = ctx.bundles[0]
            variables = ref_bundle["variables"]
            eq_functions = ref_bundle["eq_functions"]
            scalar_ineq_functions = ref_bundle["scalar_ineq_functions"]
            soc_constraints = ref_bundle["soc_constraints"]
            exp_cones = ref_bundle["exp_cones"]
            psd_cones = ref_bundle["psd_cones"]

            sol_numpy = [np.empty((B,) + v.shape, dtype=float) for v in variables]
            eq_dual = [np.empty((B,) + f.shape, dtype=float) for f in eq_functions]
            scalar_ineq_dual = [np.empty((B,) + g.shape, dtype=float) for g in scalar_ineq_functions]
            scalar_ineq_slack = [np.empty((B,) + g.shape, dtype=float) for g in scalar_ineq_functions]

            soc_dual_0 = [np.empty((B,) + c.dual_variables[0].shape, dtype=float) for c in soc_constraints]
            soc_dual_1 = [np.empty((B,) + c.dual_variables[1].shape, dtype=float) for c in soc_constraints]

            exp_dual = [
                [np.empty((B,) + dv.shape, dtype=float) for dv in c.dual_variables]
                for c in exp_cones
            ]
            psd_dual = [np.empty((B,) + c.dual_variables[0].shape, dtype=float) for c in psd_cones]

            pnorm_xstar = []
            pnorm_grad = []
            for _local_id in range(len(ref_bundle["pnorm_ineq_ids"])):
                pnorm_xstar.append([np.empty((B,) + v.shape, dtype=float) for v in variables])
                pnorm_grad.append([np.empty((B,) + v.shape, dtype=float) for v in variables])

            def _slice_params_torch(i: int):
                if ctx.batch:
                    return [p[i] if bs > 0 else p for p, bs in zip(params, ctx.batch_sizes)]
                return list(params)

            fwd_solver_iters = [0] * B
            _fwd_solve_times = [0.0] * B

            # ---- Direct SCS path: setup solvers for this forward pass ----
            if mt._warm_start:
                import scs as _scs
                import time as _time_mod

                # Detect if A/b changed by checking param tensor versions
                _param_versions = tuple(
                    p.data_ptr() if hasattr(p, 'data_ptr') else id(p)
                    for p in params
                )
                if mt._scs_data_hash != _param_versions:
                    # Use first problem to get the SCS data template via CVXPY
                    b0 = ctx.bundles[0]
                    prob0 = mt.problem_list[0]
                    params_0_np = _slice_params_np(0)
                    for pval, pparam in zip(params_0_np, b0["param_order"]):
                        pparam.value = pval
                    data0, _, _ = prob0.get_problem_data(
                        solver=cp.SCS, ignore_dpp=True
                    )
                    cone = {
                        'z': data0['dims'].zero,
                        'l': data0['dims'].nonneg,
                    }
                    scs_args = dict(
                        max_iters=int(mt._solver_args_fwd.get('max_iters', 2500)),
                        eps_abs=float(mt._solver_args_fwd.get('eps', 1e-6)),
                        eps_rel=float(mt._solver_args_fwd.get('eps', 1e-6)),
                        verbose=False,
                    )
                    # Discover the mapping once
                    if mt._scs_mapping is None:
                        n_vars = sum(int(np.prod(v.shape)) for v in b0["variables"])
                        n_eq = sum(int(np.prod(f.shape)) for f in b0["eq_functions"])
                        n_ineq = sum(int(np.prod(f.shape)) for f in b0["scalar_ineq_functions"])
                        # SCS x layout: [aux(n_vars), y_var(n_vars)]
                        # SCS y layout: [aux_dual(n_vars), eq_dual(n_eq), ineq_dual(n_ineq)]
                        mt._scs_mapping = {
                            'primal_slice': slice(n_vars, 2 * n_vars),
                            'eq_dual_slice': slice(n_vars, n_vars + n_eq),
                            'ineq_dual_slice': slice(n_vars + n_eq, n_vars + n_eq + n_ineq),
                            'c_p_slice': slice(n_vars, 2 * n_vars),
                            'b_eq_slice': slice(n_vars, n_vars + n_eq),
                        }

                    # Store dimensions for later use
                    mt._scs_c_dim = len(data0['c'])

                    # Build SCS solvers in parallel (SCS releases GIL)
                    scs_template = {
                        'P': data0['P'],
                        'A': data0['A'],
                        'b': data0['b'].copy(),
                        'c': data0['c'].copy(),
                    }
                    def _build_scs(j):
                        sd = {k: (v.copy() if isinstance(v, np.ndarray) else v)
                              for k, v in scs_template.items()}
                        return _scs.SCS(sd, cone, **scs_args)

                    with _limit_native_threads(1):
                        futs = [mt._executor.submit(_build_scs, j) for j in range(B)]
                        mt._scs_solvers = {j: f.result() for j, f in enumerate(futs)}
                    mt._scs_data_hash = _param_versions

                    # Mark that backward template needs to be built
                    mt._bwd_scs_template = None

                _m = mt._scs_mapping

            def _solve_one(i: int):
                import time as _time
                b = ctx.bundles[i]
                prob = mt.problem_list[i]

                params_i_np = _slice_params_np(i)
                # Always set CVXPY params (backward pass needs them)
                for pval, pparam in zip(params_i_np, b["param_order"]):
                    pparam.value = pval

                _t0 = _time.perf_counter()

                if mt._warm_start and i in mt._scs_solvers:
                    # ---- Direct SCS path ----
                    scs_solver = mt._scs_solvers[i]
                    # Update c vector with current p (puzzle encoding)
                    # p is at index 1 in params: [Q, p, G, h, A, b]
                    p_np = params_i_np[1]  # the puzzle-specific parameter
                    new_c = np.zeros(mt._scs_c_dim, dtype=float)
                    new_c[_m['c_p_slice']] = p_np
                    scs_solver.update(c=new_c)

                    # Warm start from cached SCS solution (same puzzle, prev epoch)
                    ws_key = mt._ws_keys[i] if mt._ws_keys is not None else None
                    ws_cached = mt._ws_cache_fwd.get(ws_key) if ws_key is not None else None
                    if ws_cached is not None and 'scs_x' in ws_cached:
                        sol = scs_solver.solve(
                            warm_start=True,
                            x=ws_cached['scs_x'],
                            y=ws_cached['scs_y'],
                            s=ws_cached['scs_s'],
                        )
                    else:
                        sol = scs_solver.solve(warm_start=False)

                    if sol['info']['status'] == 'solved' or sol['info']['status'] == 'solved_inaccurate':
                        x_scs = sol['x']
                        y_scs = sol['y']

                        # Cache full SCS state for warm starting next epoch
                        if ws_key is not None:
                            mt._ws_cache_fwd[ws_key] = {
                                'scs_x': sol['x'].copy(),
                                'scs_y': sol['y'].copy(),
                                'scs_s': sol['s'].copy(),
                            }

                        # Extract primal solution
                        y_var = x_scs[_m['primal_slice']]
                        for v_id, v in enumerate(b["variables"]):
                            vshape = v.shape
                            n_el = int(np.prod(vshape))
                            sol_numpy[v_id][i, ...] = y_var[:n_el].reshape(vshape)

                        # Extract dual values
                        eq_d = y_scs[_m['eq_dual_slice']]
                        offset = 0
                        for c_id, f in enumerate(b["eq_functions"]):
                            n_el = int(np.prod(f.shape))
                            eq_dual[c_id][i, ...] = eq_d[offset:offset+n_el].reshape(f.shape)
                            offset += n_el

                        ineq_d = y_scs[_m['ineq_dual_slice']]
                        offset = 0
                        for j, g_expr in enumerate(b["scalar_ineq_functions"]):
                            n_el = int(np.prod(g_expr.shape))
                            scalar_ineq_dual[j][i, ...] = ineq_d[offset:offset+n_el].reshape(g_expr.shape)
                            # slack = max(-(G@y - h), 0) = max(y_var, 0) for G=-I, h=0
                            scalar_ineq_slack[j][i, ...] = np.maximum(y_var[offset:offset+n_el].reshape(g_expr.shape), 0.0)
                            offset += n_el

                        fwd_solver_iters[i] = sol['info']['iter']
                        _fwd_solve_times[i] = _time.perf_counter() - _t0
                        return  # skip CVXPY path
                    else:
                        print(f"[forward] SCS direct failed for problem {i}: {sol['info']['status']}, falling back to CVXPY")

                # ---- CVXPY fallback path ----
                try:
                    prob.solve(**mt._solver_args_fwd)
                except Exception as e:
                    print(f"[forward] problem {i} solve failed: {e!r}")
                    try:
                        prob.solve(solver=cp.OSQP, warm_start=False, verbose=False)
                    except Exception as e2:
                        raise RuntimeError(f"[forward] problem {i} solve failed: {e!r} {e2!r}")
                _fwd_solve_times[i] = _time.perf_counter() - _t0

                if prob.status not in (cp.OPTIMAL, cp.OPTIMAL_INACCURATE):
                    raise RuntimeError(f"[forward] problem {i} status: {prob.status}")

                if prob.solver_stats is not None:
                    fwd_solver_iters[i] = getattr(prob.solver_stats, 'num_iters', 0)

                for v_id, v in enumerate(b["variables"]):
                    sol_numpy[v_id][i, ...] = v.value

                for c_id, c in enumerate(b["eq_constraints"]):
                    eq_dual[c_id][i, ...] = c.dual_value

                for j, g_expr in enumerate(b["scalar_ineq_functions"]):
                    g_val = np.asarray(g_expr.value, dtype=float)
                    scalar_ineq_dual[j][i, ...] = b["scalar_ineq_constraints"][j].dual_value
                    scalar_ineq_slack[j][i, ...] = np.maximum(-g_val, 0.0)

                for c_id, c in enumerate(b["soc_constraints"]):
                    dv0, dv1 = c.dual_value
                    soc_dual_0[c_id][i, ...] = dv0
                    if hasattr(dv1, "shape") and len(dv1.shape) == 2 and dv1.shape[1] == 1:
                        soc_dual_1[c_id][i, ...] = dv1.reshape(-1)
                    else:
                        soc_dual_1[c_id][i, ...] = dv1

                for c_id, c in enumerate(b["exp_cones"]):
                    shapes3 = [dv.shape for dv in c.dual_variables]
                    dv3 = _split_expcone_dual_value(c.dual_value, shapes3)
                    for k in range(3):
                        exp_dual[c_id][k][i, ...] = dv3[k]

                for c_id, c in enumerate(b["psd_cones"]):
                    psd_dual[c_id][i, ...] = c.dual_value

                if len(b["pnorm_ineq_ids"]) > 0:
                    with torch.enable_grad():
                        vars_star_t = [
                            torch.tensor(sol_numpy[v_id][i, ...], dtype=ctx.dtype, device=ctx.device, requires_grad=True)
                            for v_id in range(len(variables))
                        ]
                        params_i_det = [t.detach() for t in _slice_params_torch(i)]
                        for local_id in range(len(b["pnorm_ineq_ids"])):
                            g_t = b["pnorm_g_torch"][local_id](*vars_star_t, *params_i_det).reshape(())
                            grads = torch.autograd.grad(
                                g_t,
                                vars_star_t,
                                retain_graph=False,
                                create_graph=False,
                                allow_unused=True,
                            )
                            for v_id, gv in enumerate(grads):
                                pnorm_xstar[local_id][v_id][i, ...] = to_numpy(vars_star_t[v_id].detach())
                                pnorm_grad[local_id][v_id][i, ...] = 0.0 if gv is None else to_numpy(gv.detach())

            with _limit_native_threads(1):
                futs = [mt._executor.submit(_solve_one, i) for i in range(B)]
                for f in futs:
                    f.result()

                # single thread for debugging
                # for i in range(B):
                #     _solve_one(i) 

            if mt._warm_start:
                total_fwd = sum(fwd_solver_iters)
                max_solve = max(_fwd_solve_times)
                sum_solve = sum(_fwd_solve_times)
                fwd_cache_size = len(mt._ws_cache_fwd)
                print(f"[forward] iters: avg={total_fwd/max(B,1):.0f}, max_solve={max_solve:.3f}s, sum_solve={sum_solve:.3f}s, cache={fwd_cache_size}")
            elif mt.verbose:
                total_fwd = sum(fwd_solver_iters)
                print(f"[forward] solver iters: total={total_fwd}, avg={total_fwd/max(B,1):.1f}")

            ctx.sol_numpy = sol_numpy
            ctx.eq_dual = eq_dual
            ctx.scalar_ineq_dual = scalar_ineq_dual
            ctx.scalar_ineq_slack = scalar_ineq_slack
            ctx.soc_dual_0 = soc_dual_0
            ctx.soc_dual_1 = soc_dual_1
            ctx.exp_dual = exp_dual
            ctx.psd_dual = psd_dual
            ctx.pnorm_xstar = pnorm_xstar
            ctx.pnorm_grad = pnorm_grad
            ctx.params = params

            # if want to check active counts
            if mt.verbose:
                ctx.active_counts = active_counts_dict(ctx)
                print(f"active_counts: {ctx.active_counts}")

            sol_torch = [to_torch(arr, ctx.dtype, ctx.device) for arr in sol_numpy]
            return tuple(sol_torch) # return the solution

        @staticmethod
        def backward(ctx, *dvars):
            mt = ctx.mt
            bundles = ctx.bundles
            B = ctx.batch_size

            ref = bundles[0]
            num_vars = len(ref["variables"])
            num_scalar_ineq = len(ref["scalar_ineq_functions"])

            params_np_all = [to_numpy(p) for p in ctx.params]
            dvars_np_all = [to_numpy(dv) for dv in dvars]

            def _slice_params_np(i: int):
                if ctx.batch:
                    return [arr[i] if bs > 0 else arr for arr, bs in zip(params_np_all, ctx.batch_sizes)]
                return params_np_all

            # def _slice_dvars_np(i: int):
            #     out = []
            #     for arr, v in zip(dvars_np_all, ref["variables"]):
            #         vshape = tuple(v.shape)
            #         if arr.shape == (B,) + vshape:
            #             out.append(arr[i])
            #         elif B == 1 and arr.ndim >= 1 and arr.shape[0] == 1 and tuple(arr.shape[1:]) == vshape:
            #             out.append(arr[0])
            #         else:
            #             out.append(arr)
            #     return out
            def _slice_dvars_np(i: int):
                if ctx.batch:
                    return [arr[i] for arr in dvars_np_all]
                # Even when not batched, sol_numpy has leading B=1 dim,
                # so dvars also has shape (1, *v.shape). Slice it out.
                return [arr[i] if arr.ndim > len(v.shape) else arr
                        for arr, v in zip(dvars_np_all, ref["variables"])]

            y_dim = int(np.prod((_slice_dvars_np(0)[0]).shape))
            num_eq = int(np.prod(ctx.eq_dual[0][0].shape)) if (len(ctx.eq_dual) > 0 and ctx.batch) else (
                int(np.prod(ctx.eq_dual[0].shape)) if len(ctx.eq_dual) > 0 else 0
            )
            cap_scalar = int(max(1, y_dim - num_eq))

            new_sol_lagrangian = [np.empty_like(ctx.sol_numpy[k]) for k in range(num_vars)]
            new_eq_dual = [np.empty_like(ctx.eq_dual[k]) for k in range(len(ref["eq_constraints"]))]

            new_active_dual = [np.empty((B,) + c.shape, dtype=float) for c in ref["active_eq_constraints"]]

            new_soc_lam = [np.zeros((B,), dtype=float) for _ in ref["soc_lin_constraints"]]
            new_pnorm_lam = [np.zeros((B,), dtype=float) for _ in ref["pnorm_tangent_constraints"]]

            new_exp_dual = [
                [np.empty_like(ctx.exp_dual[j][k]) for k in range(3)]
                for j in range(len(ref["exp_cones"]))
            ]
            new_psd_dual = [np.empty_like(ctx.psd_dual[k]) for k in range(len(ref["psd_cones"]))]

            def _slice_params_torch(i: int, params_src):
                if ctx.batch:
                    return [p[i] if bs > 0 else p for p, bs in zip(params_src, ctx.batch_sizes)]
                return list(params_src)

            bwd_solver_iters = [0] * B
            _bwd_solve_times = [0.0] * B

            # ---- Rebuild backward SCS template when params change ----
            if mt._warm_start and (getattr(mt, '_bwd_scs_template', None) is None
                                    or getattr(mt, '_bwd_param_hash', None) != id(ctx.params)):
                import scipy.sparse as _sp_init
                b0 = bundles[0]
                prob0_bwd = mt.perturbed_problem_list[0]
                n_vars_total = y_dim
                n_eq_total = num_eq
                # Set params and mask=1 so all mask entries exist in A
                params_0_np = _slice_params_np(0)
                for pval, pparam in zip(params_0_np, b0["param_order"]):
                    pparam.value = pval
                for dv in b0["dvar_params"]:
                    dv.value = np.zeros(dv.shape)
                for dp in b0.get("eq_dual_params", []):
                    dp.value = np.zeros(dp.shape)
                for dp in b0.get("scalar_ineq_dual_params", []):
                    dp.value = np.zeros(dp.shape)
                for mp in b0.get("scalar_active_mask_params", []):
                    mp.value = np.ones(mp.shape)
                bwd_data_ones, _, _ = prob0_bwd.get_problem_data(
                    solver=cp.SCS, ignore_dpp=True
                )
                bwd_A_csc = _sp_init.csc_matrix(bwd_data_ones['A'])
                mask_row_start = n_vars_total + n_eq_total
                mask_data_indices = np.empty(n_vars_total, dtype=int)
                for k in range(n_vars_total):
                    col_s, col_e = bwd_A_csc.indptr[k], bwd_A_csc.indptr[k+1]
                    rows = bwd_A_csc.indices[col_s:col_e]
                    idx = np.searchsorted(rows, mask_row_start + k)
                    mask_data_indices[k] = col_s + idx
                # Get baseline c with mask=0
                for mp in b0.get("scalar_active_mask_params", []):
                    mp.value = np.zeros(mp.shape)
                bwd_data_base, _, _ = prob0_bwd.get_problem_data(
                    solver=cp.SCS, ignore_dpp=True
                )
                bwd_eps = float(mt._solver_args_bwd.get('eps', 1e-5))
                mt._bwd_scs_template = {
                    'c_base': bwd_data_base['c'].copy(),
                    'b': bwd_data_base['b'].copy(),
                    'P': bwd_data_ones['P'],
                    'A_base_data': bwd_A_csc.data.copy(),
                    'A_base_indices': bwd_A_csc.indices.copy(),
                    'A_base_indptr': bwd_A_csc.indptr.copy(),
                    'A_shape': bwd_A_csc.shape,
                    'mask_data_indices': mask_data_indices,
                    'alpha': float(b0['alpha']),
                    'cone': {
                        'z': bwd_data_ones['dims'].zero,
                        'l': bwd_data_ones['dims'].nonneg,
                    },
                    'scs_args': dict(
                        max_iters=int(mt._solver_args_bwd.get('max_iters', 2500)),
                        eps_abs=bwd_eps, eps_rel=bwd_eps, verbose=False,
                    ),
                }
                mt._bwd_param_hash = id(ctx.params)

            def _solve_perturbed_one(i: int):
                import time as _time
                b = bundles[i]
                prob = mt.perturbed_problem_list[i]

                params_i_np = _slice_params_np(i)
                for pval, pparam in zip(params_i_np, b["param_order"]):
                    pparam.value = pval

                dvals_i = _slice_dvars_np(i)
                for j, v in enumerate(b["variables"]):
                    b["dvar_params"][j].value = dvals_i[j]
                    v.value = ctx.sol_numpy[j][i, ...]

                for j in range(len(b["eq_functions"])):
                    b["eq_dual_params"][j].value = ctx.eq_dual[j][i]

                cap = cap_scalar

                scalar_candidates = []
                for j in b["scalar_scalar_indices"]:
                    sl_s = float(np.asarray(ctx.scalar_ineq_slack[j][i]).reshape(()))
                    lam_s = float(np.asarray(ctx.scalar_ineq_dual[j][i]).reshape(()))
                    lam_s = 0.0 if lam_s < -1e-8 else max(lam_s, 0.0)
                    if sl_s <= mt.slack_tol and lam_s >= mt.dual_cutoff:
                        scalar_candidates.append((lam_s, j))

                if len(scalar_candidates) > 0:
                    scalar_candidates.sort(key=lambda t: t[0])
                    active_scalar = set([j for _, j in scalar_candidates[-cap:]]) if len(scalar_candidates) > cap else set([j for _, j in scalar_candidates])
                else:
                    active_scalar = set()

                for j in range(num_scalar_ineq):
                    lam = np.asarray(ctx.scalar_ineq_dual[j][i], dtype=float)
                    lam = np.where(lam < -1e-8, lam, np.maximum(lam, 0.0))
                    b["scalar_ineq_dual_params"][j].value = lam

                    gshape = b["scalar_ineq_functions"][j].shape
                    if int(np.prod(gshape)) == 1:
                        b["scalar_active_mask_params"][j].value = 1.0 if (j in active_scalar) else 0.0
                    else:
                        sl = np.asarray(ctx.scalar_ineq_slack[j][i], dtype=float)
                        mask = (sl <= mt.slack_tol).astype(np.float64)
                        cap_vec = cap_scalar
                        if mask.sum() > cap_vec:
                            lam_flat = lam.reshape(-1)
                            idx = np.argpartition(lam_flat, -cap_vec)[-cap_vec:]
                            mask_flat = np.zeros_like(lam_flat, dtype=np.float64)
                            mask_flat[idx] = 1.0
                            mask = mask_flat.reshape(lam.shape)
                        b["scalar_active_mask_params"][j].value = mask

                for j in range(len(b["soc_constraints"])):
                    b["soc_dual_params_0"][j].value = np.maximum(ctx.soc_dual_0[j][i], 0.0)
                    b["soc_dual_params_1"][j].value = ctx.soc_dual_1[j][i]

                for j in range(len(b["exp_cones"])):
                    for k in range(3):
                        b["exp_dual_params"][j][k].value = ctx.exp_dual[j][k][i]

                for j in range(len(b["psd_cones"])):
                    b["psd_dual_params"][j].value = ctx.psd_dual[j][i]

                for local_id, j_scalar in enumerate(b["pnorm_ineq_ids"]):
                    for v_id in range(num_vars):
                        b["pnorm_xstar_params"][local_id][v_id].value = ctx.pnorm_xstar[local_id][v_id][i]
                        b["pnorm_grad_params"][local_id][v_id].value = ctx.pnorm_grad[local_id][v_id][i]

                _bwd_t0 = _time.perf_counter()
                if mt._warm_start and hasattr(mt, '_bwd_scs_template'):
                    # ---- Direct SCS path for backward (no CVXPY) ----
                    import scs as _scs_bwd
                    import scipy.sparse as _sp
                    try:
                        tmpl = mt._bwd_scs_template

                        # Backward SCS variable order: x = [y, t] (opposite of forward)
                        # c[0:y_dim] = p + dvar/alpha - lambda (all y-linear terms)
                        # c[y_dim:2*y_dim] = 0 (t has no linear cost)
                        new_c = tmpl['c_base'].copy()
                        dvar_i = dvals_i[0].ravel()
                        lam_i = np.asarray(b["scalar_ineq_dual_params"][0].value, dtype=float).ravel()
                        p_i = params_i_np[1].ravel()
                        new_c[:y_dim] = p_i + dvar_i / tmpl['alpha'] - lam_i

                        # Construct A: copy base, update mask diagonal
                        A_data = tmpl['A_base_data'].copy()
                        mask_i = np.asarray(b["scalar_active_mask_params"][0].value, dtype=float).ravel()
                        mask_data_idx = tmpl['mask_data_indices']
                        for k in range(y_dim):
                            A_data[mask_data_idx[k]] = -mask_i[k]
                        A_sparse = _sp.csc_matrix(
                            (A_data, tmpl['A_base_indices'], tmpl['A_base_indptr']),
                            shape=tmpl['A_shape'],
                        )

                        # Update b: b[y_dim:y_dim+num_eq] = b_eq (index 5 in params)
                        # b vector matches the A rows: [aux(y_dim), eq(num_eq), mask(y_dim)]
                        new_b = tmpl['b'].copy()
                        b_eq_i = params_i_np[5].ravel()
                        new_b[y_dim:y_dim + num_eq] = b_eq_i  # same position confirmed earlier

                        bwd_sd = {
                            'P': tmpl['P'], 'A': A_sparse,
                            'b': new_b, 'c': new_c,
                        }
                        bwd_solver = _scs_bwd.SCS(bwd_sd, tmpl['cone'], **tmpl['scs_args'])

                        # Warm start backward from cached SCS state
                        bwd_ws_key = mt._ws_keys[i] if mt._ws_keys is not None else None
                        bwd_ws = mt._ws_cache_bwd.get(bwd_ws_key) if bwd_ws_key is not None else None
                        if bwd_ws is not None:
                            bwd_sol = bwd_solver.solve(
                                warm_start=True,
                                x=bwd_ws['scs_x'], y=bwd_ws['scs_y'], s=bwd_ws['scs_s'],
                            )
                        else:
                            bwd_sol = bwd_solver.solve(warm_start=False)

                        if bwd_sol['info']['status'] in ('solved', 'solved_inaccurate'):
                            x_bwd = bwd_sol['x']
                            y_bwd = bwd_sol['y']
                            bwd_solver_iters[i] = bwd_sol['info']['iter']

                            y_var = x_bwd[:y_dim]
                            for j, v in enumerate(b["variables"]):
                                vshape = v.shape
                                n_el = int(np.prod(vshape))
                                new_sol_lagrangian[j][i, ...] = y_var[:n_el].reshape(vshape)

                            eq_d = y_bwd[y_dim:y_dim + num_eq]
                            offset = 0
                            for c_id, f in enumerate(b["eq_functions"]):
                                n_el = int(np.prod(f.shape))
                                new_eq_dual[c_id][i, ...] = eq_d[offset:offset+n_el].reshape(f.shape)
                                offset += n_el

                            active_d = y_bwd[y_dim + num_eq:]
                            offset = 0
                            for c_id, c_expr in enumerate(b["active_eq_constraints"]):
                                n_el = int(np.prod(c_expr.shape))
                                new_active_dual[c_id][i, ...] = active_d[offset:offset+n_el].reshape(c_expr.shape)
                                offset += n_el

                            # Cache backward SCS state
                            if bwd_ws_key is not None:
                                mt._ws_cache_bwd[bwd_ws_key] = {
                                    'scs_x': bwd_sol['x'].copy(),
                                    'scs_y': bwd_sol['y'].copy(),
                                    'scs_s': bwd_sol['s'].copy(),
                                }

                            _bwd_solve_times[i] = _time.perf_counter() - _bwd_t0
                            return
                        else:
                            print(f"[backward] SCS direct failed for {i}: {bwd_sol['info']['status']}, fallback")
                    except Exception as e:
                        print(f"[backward] SCS direct error for {i}: {e!r}, fallback")

                # ---- CVXPY fallback path ----
                try:
                    prob.solve(**mt._solver_args_bwd)
                except Exception as e:
                    print(f"[backward] problem {i} perturbed solve failed: {e!r}")
                    try:
                        b["perturbed_problem"].solve(solver=cp.OSQP, eps_abs=1e-4, eps_rel=1e-4, warm_start=True, verbose=False)
                    except Exception as e2:
                        raise RuntimeError(f"[backward] problem {i} perturbed solve failed: {e!r} {e2!r}")

                if prob.status not in (cp.OPTIMAL, cp.OPTIMAL_INACCURATE):
                    raise RuntimeError(f"[backward] perturbed problem {i} status: {prob.status}")

                if prob.solver_stats is not None:
                    bwd_solver_iters[i] = getattr(prob.solver_stats, 'num_iters', 0)

                for j, v in enumerate(b["variables"]):
                    new_sol_lagrangian[j][i, ...] = v.value

                for c_id, c in enumerate(b["eq_constraints"]):
                    new_eq_dual[c_id][i, ...] = c.dual_value

                for c_id, c in enumerate(b["active_eq_constraints"]):
                    new_active_dual[c_id][i, ...] = c.dual_value

                for c_id, c in enumerate(b["soc_lin_constraints"]):
                    dv = c.dual_value
                    new_soc_lam[c_id][i] = 0.0 if dv is None else float(np.asarray(dv).reshape(()))

                for c_id, c in enumerate(b["pnorm_tangent_constraints"]):
                    dv = c.dual_value
                    lam_val = 0.0 if dv is None else float(np.asarray(dv).reshape(()))
                    j_scalar = b["pnorm_ineq_ids"][c_id]
                    mval = float(np.asarray(b["scalar_active_mask_params"][j_scalar].value).reshape(()))
                    if mval < 0.5:
                        lam_val = 0.0
                    new_pnorm_lam[c_id][i] = lam_val

                for c_id, c in enumerate(b["exp_cones"]):
                    shapes3 = [dv.shape for dv in c.dual_variables]
                    dv3 = _split_expcone_dual_value(c.dual_value, shapes3)
                    for k in range(3):
                        new_exp_dual[c_id][k][i, ...] = dv3[k]

                for c_id, c in enumerate(b["psd_cones"]):
                    new_psd_dual[c_id][i, ...] = c.dual_value

            with _limit_native_threads(1):
                futs = [mt._executor.submit(_solve_perturbed_one, i) for i in range(B)]
                for f in futs:
                    f.result()

            if mt._warm_start or mt.verbose:
                total_bwd = sum(bwd_solver_iters)
                bwd_max = max(_bwd_solve_times)
                bwd_sum = sum(_bwd_solve_times)
                print(f"[backward] iters: avg={total_bwd/max(B,1):.0f}, max_solve={bwd_max:.3f}s, sum_solve={bwd_sum:.3f}s")

            new_sol = [to_torch(v, ctx.dtype, ctx.device) for v in new_sol_lagrangian]
            vars_old = [to_torch(ctx.sol_numpy[j], ctx.dtype, ctx.device) for j in range(num_vars)]

            new_eq_dual_t = [to_torch(v, ctx.dtype, ctx.device) for v in new_eq_dual]
            old_eq_dual_t = [to_torch(v, ctx.dtype, ctx.device) for v in ctx.eq_dual]

            old_scalar_dual_t = [to_torch(v, ctx.dtype, ctx.device) for v in ctx.scalar_ineq_dual]
            new_active_dual_t = [to_torch(v, ctx.dtype, ctx.device) for v in new_active_dual]

            new_exp_dual_t = [
                [to_torch(new_exp_dual[j][k], ctx.dtype, ctx.device) for k in range(3)]
                for j in range(len(ref["exp_cones"]))
            ]
            new_psd_dual_t = [to_torch(v, ctx.dtype, ctx.device) for v in new_psd_dual]

            new_pnorm_lam_t = [to_torch(v, ctx.dtype, ctx.device) for v in new_pnorm_lam]

            params_req = []
            req_grad_mask = []
            for p in ctx.params:
                need = bool(getattr(p, "requires_grad", False))
                q = p.detach()
                if need:
                    q.requires_grad_(True)
                params_req.append(q)
                req_grad_mask.append(need)

            loss = 0.0
            with torch.enable_grad():
                for i in range(B):
                    b = bundles[i]
                    vars_new_i = [v[i] for v in new_sol]
                    vars_old_i = [v[i] for v in vars_old]
                    params_i = slice_params_for_batch(params_req, ctx.batch_sizes, i) if ctx.batch else params_req

                    new_eq_dual_i = [d[i] for d in new_eq_dual_t]
                    old_eq_dual_i = [d[i] for d in old_eq_dual_t]

                    new_scalar_dual_full_i = []
                    ptr = 0
                    for j in range(num_scalar_ineq):
                        if j in b["non_pnorm_set"]:
                            new_scalar_dual_full_i.append(new_active_dual_t[ptr][i])
                            ptr += 1
                        elif j in b["pnorm_set"]:
                            lid = b["pnorm_map"][j]
                            new_scalar_dual_full_i.append(new_pnorm_lam_t[lid][i])
                        else:
                            new_scalar_dual_full_i.append(old_scalar_dual_t[j][i])
                    old_scalar_dual_full_i = [d[i] for d in old_scalar_dual_t]

                    new_exp_dual_i = []
                    for j in range(len(b["exp_cones"])):
                        for k in range(3):
                            new_exp_dual_i.append(new_exp_dual_t[j][k][i])

                    new_psd_dual_i = [d[i] for d in new_psd_dual_t]

                    phi_new = b["phi_torch"](*vars_new_i, *params_i)
                    phi_old = b["phi_torch"](*vars_old_i, *params_i)

                    eq_new = b["eq_dual_term_torch"](*vars_old_i, *params_i, *new_eq_dual_i)
                    eq_old = b["eq_dual_term_torch"](*vars_old_i, *params_i, *old_eq_dual_i)


                    ineq_new = b["ineq_dual_term_torch"](*vars_old_i, *params_i, *new_scalar_dual_full_i)
                    ineq_old = b["ineq_dual_term_torch"](*vars_old_i, *params_i, *old_scalar_dual_full_i)

                    if b["exp_dual_term_torch"] is not None:
                        exp_new = b["exp_dual_term_torch"](*vars_old_i, *params_i, *new_exp_dual_i)
                    else:
                        exp_new = 0.0

                    if b["psd_dual_term_torch"] is not None:
                        psd_new = b["psd_dual_term_torch"](*vars_old_i, *params_i, *new_psd_dual_i)
                    else:
                        psd_new = 0.0

                    loss = loss + (phi_new + ineq_new + eq_new + exp_new + psd_new - phi_old - eq_old - ineq_old)

                loss = mt.alpha * loss

            grads_req = torch.autograd.grad(
                outputs=loss,
                inputs=[q for q, need in zip(params_req, req_grad_mask) if need],
                allow_unused=True,
                retain_graph=False,
            )

            grads = []
            it = iter(grads_req)
            for need in req_grad_mask:
                grads.append(next(it) if need else None)

            return tuple(grads)
    return _FFOLayerFn