Add a CPU Moreau backend for MLX

docs/api/mlx.rst

@@ -74,15 +74,58 @@ Compute both value and gradient efficiently:
    value_and_grad_fn = mx.value_and_grad(loss_fn, argnums=[0, 1])
    loss_val, (dA, db) = value_and_grad_fn(A_mx, b_mx)
 
-Apple Silicon Optimization
---------------------------
+Moreau CPU Backend
+------------------
+
+Select Moreau explicitly to use its CPU interior-point solver and implicit
+derivatives. Install the ``mlx`` and ``moreau`` extras (on Linux, also install
+``mlx[cpu]``):
+
+.. code-block:: bash
+
+   pip install "cvxpylayers[mlx,moreau]"
+
+.. code-block:: python
+
+   import cvxpy as cp
+   import mlx.core as mx
+   from cvxpylayers.mlx import CvxpyLayer
+
+   x = cp.Variable(2)
+   q = cp.Parameter(2)
+   problem = cp.Problem(
+       cp.Minimize(0.5 * cp.sum_squares(x) + q @ x),
+       [x >= 0],
+   )
+   layer = CvxpyLayer(problem, parameters=[q], variables=[x], solver=cp.MOREAU)
+
+   # This stream also keeps the surrounding MLX operations on CPU.
+   with mx.stream(mx.cpu):
+       q_mx = mx.array([-2.0, 3.0])
+       x_sol = layer(q_mx)[0]  # [2.0, 0.0]
+       dq = mx.grad(lambda q: mx.sum(layer(q)[0]))(q_mx)  # [-1.0, 0.0]
+       mx.eval(x_sol, dq)
+
+Moreau's solve and backward pass always run on CPU, including when MLX's default
+device is a GPU. This interface does not accelerate Moreau with Metal or CUDA.
+It transfers canonical problem data to NumPy in double precision and returns MLX
+arrays in the input dtype. PyTorch and JAX are not required.
+
+The backend supports quadratic objectives, batched and broadcast parameters,
+primal and dual outputs, and ``mx.grad`` / ``mx.value_and_grad``. Each forward
+call retains its own solver state until its backward pass. It uses Moreau's IPM
+algorithm; explicit GPU devices and the active-set algorithm are rejected.
+Solver settings can be provided at construction or overridden per call, for example
+``solver_args={"max_iter": 100}``.
+
+MLX Device Selection
+--------------------
 
-MLX automatically uses the unified memory architecture of Apple Silicon,
-allowing efficient computation without explicit device management:
+MLX operations surrounding the solver follow MLX's selected stream. On Apple
+Silicon they can use unified memory, while the optimization solve runs on CPU:
 
 .. code-block:: python
 
-   # No need to move tensors to GPU - MLX handles this automatically
    A_mx = mx.random.normal((1000, 500))
    b_mx = mx.random.normal((1000,))
 
@@ -94,4 +137,4 @@ Notes
 
 - MLX uses lazy evaluation; call ``mx.eval()`` to force computation
 - The MLX layer supports batched execution like PyTorch and JAX
-- For best performance on Apple Silicon, prefer MLX over PyTorch CPU
+- Moreau's MLX backend uses CPU solves and first-order reverse-mode derivatives

docs/guide/solvers.md

@@ -175,6 +175,10 @@ layer = CvxpyLayer(problem, parameters=[A, b], variables=[x],
 [Moreau](https://docs.moreau.so/) is the default backend for PyTorch and JAX on CPU and GPU.
 It supports PyTorch and JAX with native autograd integration, warm starts, and `jax.jit` compatibility.
 
+MLX also supports Moreau via `solver=cp.MOREAU`. Its solve and backward pass use
+the CPU interior-point solver, without Metal or CUDA acceleration. See the
+[MLX API](../api/mlx.rst) for installation and a gradient example.
+
 ### Setup
 
 Moreau >= 0.4.0 is installed with CVXPYlayers. It is open source under the Apache

src/cvxpylayers/interfaces/moreau_if.py

@@ -7,9 +7,10 @@
 
 where K is a product of cones.
 
-Uses Moreau's native PyTorch and JAX solvers with built-in automatic differentiation:
+Uses Moreau's native solvers for forward solves and implicit differentiation:
 - PyTorch: moreau.torch.Solver with two-step API (setup + solve) and autograd support
 - JAX: moreau.jax.Solver with custom_vjp for gradients
+- MLX: moreau.CompiledSolver on CPU with the layer's custom VJP
 
 Limitations:
 - Not thread-safe: solver instances are lazily initialized and cached on MOREAU_ctx.
@@ -49,6 +50,11 @@
 except ImportError:
     torch = None  # type: ignore[assignment]
 
+try:
+    import mlx.core as mx
+except ImportError:
+    mx = None  # type: ignore[assignment]
+
 if TYPE_CHECKING:
     import cvxpylayers.utils.parse_args as pa
 
@@ -292,11 +298,12 @@ def cones(self):
             )
         return self._cones
 
-    def _get_settings(self, enable_grad: bool = True):
+    def _get_settings(self, enable_grad: bool = True, overrides=None):
         """Get moreau.Settings configured from self.options.
 
         Args:
             enable_grad: Whether to enable gradient computation in Moreau.
+            overrides: Per-call settings that override the layer defaults.
 
         Accepts any moreau.Settings field names directly (e.g., max_iter,
         tol_gap_abs, verbose, etc.).  ``ipm_settings`` can be a dict or a
@@ -305,7 +312,7 @@ def _get_settings(self, enable_grad: bool = True):
         settings = moreau.Settings(enable_grad=enable_grad)
 
         # Set any field that exists on moreau.Settings
-        for key, value in self.options.items():
+        for key, value in {**self.options, **(overrides or {})}.items():
             if key == "ipm_settings":
                 if isinstance(value, dict):
                     value = moreau.IPMSettings(**value)
@@ -396,6 +403,36 @@ def get_jax_solver(self):
             self._jax_solver = self._create_jax_solver()
         return self._jax_solver
 
+    def mlx_to_data(self, quad_obj_values, lin_obj_values, con_values) -> "MOREAU_mlx_data":
+        """Copy MLX canonical data to NumPy for Moreau's CPU backend."""
+        if mx is None:
+            raise ImportError("MLX interface requires 'mlx'. Install with: pip install mlx")
+        if moreau is None:
+            raise ImportError("Moreau solver requires 'moreau'. Install with: pip install moreau")
+
+        batch_size, originally_unbatched = _detect_batch_size(con_values)
+        # Parametrization matrices already produce CSR-ordered values. Moreau
+        # expects float64 arrays with the batch dimension first.
+        con_np = np.array(con_values, dtype=np.float64).reshape(-1, batch_size).T
+        lin_np = np.array(lin_obj_values, dtype=np.float64).reshape(-1, batch_size).T
+        P_values = (
+            np.array(quad_obj_values, dtype=np.float64).reshape(-1, batch_size).T
+            if quad_obj_values is not None
+            else np.zeros((batch_size, 0))
+        )
+        b = np.zeros((batch_size, self.A_shape[0]))
+        b[:, self.b_idx] = con_np[:, self.nnz_A :]
+
+        return MOREAU_mlx_data(
+            ctx=self,
+            P_values=P_values,
+            A_values=-con_np[:, : self.nnz_A],
+            q=lin_np[:, :-1],
+            b=b,
+            originally_unbatched=originally_unbatched,
+            dtype=lin_obj_values.dtype,
+        )
+
     def torch_to_data(self, quad_obj_values, lin_obj_values, con_values) -> "MOREAU_data":
         """Prepare data for torch solve.
 
@@ -479,6 +516,79 @@ def torch_to_data(self, quad_obj_values, lin_obj_values, con_values) -> "MOREAU_
         )
 
 
+@dataclass
+class MOREAU_mlx_data:
+    """CPU solve data; each forward call retains its own solver for the VJP."""
+
+    ctx: MOREAU_ctx
+    P_values: np.ndarray
+    A_values: np.ndarray
+    q: np.ndarray
+    b: np.ndarray
+    originally_unbatched: bool
+    dtype: Any
+
+    def mlx_solve(self, solver_args=None):
+        settings = self.ctx._get_settings(overrides=solver_args)
+        if settings.device not in ("auto", "cpu"):
+            raise ValueError("The Moreau MLX backend only supports device='cpu'.")
+        settings.device = "cpu"
+        if settings.solver not in ("auto", "ipm"):
+            raise ValueError("The Moreau MLX backend requires solver='ipm' for differentiation.")
+        settings.solver = "ipm"
+        settings.enable_grad = True
+        settings.batch_size = self.q.shape[0]
+
+        # CompiledSolver.backward uses the last solve's state. A solver owned by
+        # this forward call keeps multiple calls to the same layer independent.
+        solver = moreau.CompiledSolver(
+            n=self.ctx.P_shape[0],
+            m=self.ctx.A_shape[0],
+            P_row_offsets=self.ctx.P_row_offsets,
+            P_col_indices=self.ctx.P_col_indices,
+            A_row_offsets=self.ctx.A_row_offsets,
+            A_col_indices=self.ctx.A_col_indices,
+            cones=self.ctx.cones,
+            settings=settings,
+        )
+        solver.setup(self.P_values, self.A_values)
+        solution = solver.solve(self.q, self.b)
+        for i, status in enumerate(solver.info.status):
+            if status not in (moreau.SolverStatus.Solved, moreau.SolverStatus.AlmostSolved):
+                raise RuntimeError(f"Moreau failed for batch element {i}: {status}")
+        dual = (
+            np.concatenate((solution.z, solution.z_x), axis=-1)
+            if self.ctx.dir_cones
+            else solution.z
+        )
+        return (
+            mx.array(solution.x * self.ctx.primal_scale, dtype=self.dtype),
+            mx.array(dual, dtype=self.dtype),
+            solver,
+        )
+
+    def mlx_derivative(self, dprimal, ddual, solver):
+        ddual = np.array(ddual, dtype=np.float64)
+        m = self.ctx.A_shape[0]
+        grads = solver.backward(
+            np.array(dprimal, dtype=np.float64) * self.ctx.primal_scale,
+            ddual[:, :m],
+            dz_x=ddual[:, m:] if self.ctx.dir_cones else None,
+        )
+        # Undo q's constant objective offset and the A sign / sparse b scatter.
+        dP = grads["dP_values"].T if self.ctx.nnz_P else None
+        dq = np.concatenate((grads["dq"], np.zeros((self.q.shape[0], 1))), axis=1).T
+        dA = np.concatenate((-grads["dA_values"], grads["db"][:, self.ctx.b_idx]), axis=1).T
+        if self.originally_unbatched:
+            dP = dP[:, 0] if dP is not None else None
+            dq, dA = dq[:, 0], dA[:, 0]
+        return (
+            mx.array(dP, dtype=self.dtype) if dP is not None else None,
+            mx.array(dq, dtype=self.dtype),
+            mx.array(dA, dtype=self.dtype),
+        )
+
+
 @dataclass
 class MOREAU_data:
     """Data class for PyTorch Moreau solver.

src/cvxpylayers/mlx/cvxpylayer.py

@@ -63,10 +63,7 @@ def _apply_gp_log_transform(
         return params
 
     # Use pre-computed mask for consistency with Torch/JAX (no dict lookups)
-    return tuple(
-        mx.log(p) if needs_log else p
-        for p, needs_log in zip(params, ctx.gp_log_mask)
-    )
+    return tuple(mx.log(p) if needs_log else p for p, needs_log in zip(params, ctx.gp_log_mask))
 
 
 def _flatten_and_batch_params(
@@ -186,7 +183,7 @@ def _unpack_primal_svec(svec: mx.array, n: int, batch: tuple) -> mx.array:
     return _svec_to_symmetric(svec, n, batch, rows, cols)
 
 
-def _unpack_svec(svec: mx.array, n: int, batch: tuple) -> mx.array:
+def _unpack_svec(svec: mx.array, n: int, batch: tuple, upper: bool = False) -> mx.array:
     """Unpack scaled vectorized (svec) form to full symmetric matrix.
 
     The svec format stores a symmetric n x n matrix as a vector of length n*(n+1)/2,
@@ -197,11 +194,12 @@ def _unpack_svec(svec: mx.array, n: int, batch: tuple) -> mx.array:
         svec: Scaled vectorized form
         n: Matrix dimension
         batch: Batch dimensions
+        upper: Whether the solver uses upper-triangular column-major svec ordering.
 
     Returns:
         Full symmetric matrix with scaling removed
     """
-    rows_rm, cols_rm = np.tril_indices(n)
+    rows_rm, cols_rm = np.triu_indices(n) if upper else np.tril_indices(n)
     sort_idx = np.lexsort((rows_rm, cols_rm))
     rows = rows_rm[sort_idx]
     cols = cols_rm[sort_idx]
@@ -240,12 +238,14 @@ def _recover_results(
             data = primal[..., var.primal]
         else:  # var.source == "dual"
             data = dual[..., var.dual]
+            if var.dual_indices is not None:
+                data = mx.take(dual, mx.array(var.dual_indices), axis=-1)
 
         # Use pre-computed unpack_fn field (consistent with Torch/JAX)
         if var.unpack_fn == "svec_primal":
             results.append(_unpack_primal_svec(data, var.shape[0], batch_shape))
         elif var.unpack_fn == "svec_dual":
-            results.append(_unpack_svec(data, var.shape[0], batch_shape))
+            results.append(_unpack_svec(data, var.shape[0], batch_shape, upper=var.dual_upper))
         elif var.unpack_fn == "reshape":
             results.append(_reshape_fortran(data, batch_shape + var.shape))
         else:
@@ -256,8 +256,7 @@ def _recover_results(
     # Uses pre-computed source field (consistent with Torch/JAX)
     if ctx.gp:
         results = [
-            mx.exp(r) if var.source == "primal" else r
-            for r, var in zip(results, ctx.var_recover)
+            mx.exp(r) if var.source == "primal" else r for r, var in zip(results, ctx.var_recover)
         ]
 
     # Squeeze batch dimension for unbatched inputs
@@ -321,7 +320,8 @@ def __init__(
                 at runtime. Order must match the order of arrays passed to __call__().
             variables: List of CVXPY Variables whose optimal values will be returned
                 by __call__(). Order determines the order of returned arrays.
-            solver: CVXPY solver to use (e.g., ``cp.CLARABEL``, ``cp.SCS``).
+            solver: Solver backend to use (``cp.DIFFCP`` or ``cp.MOREAU``).
+                Moreau solves and differentiates on CPU using its IPM solver.
                 If None, uses the default diffcp solver.
             gp: If True, problem is a geometric program. Parameters will be
                 log-transformed before solving.