Add cp.stack atom (cvxpy#2956)

* Add cp.stack atom

* rebase

* address comments: remove erroneous import, add canonicalization test

* Run tests against editable install

* Enable editable install for OpenMP builds

* Use editable install in backend workflow

* Fix Windows install mode

* Revert Windows-specific pytest args

Author: ClayCampaigne

.github/workflows/test_backends.yml

@@ -18,12 +18,11 @@ jobs:
       - uses: actions/checkout@v5
       - name: Install cvxpy dependencies
         run: |
-          pip install .
+          pip install -e .
           pip install pytest hypothesis
       - name: Run tests for each non-default backend
         run : |
           export CVXPY_DEFAULT_CANON_BACKEND="SCIPY"
           python -c "from cvxpy.cvxcore.python.canonInterface import get_default_canon_backend; print(get_default_canon_backend())"
           python -c "from cvxpy.cvxcore.python.canonInterface import get_default_canon_backend; assert get_default_canon_backend() == 'SCIPY'"
           pytest
-

continuous_integration/test_script.sh

@@ -16,11 +16,11 @@ python -c "import numpy; print('numpy %s' % numpy.__version__)"
 python -c "import scipy; print('scipy %s' % scipy.__version__)"
 
 if [ $USE_OPENMP == "True" ] && [ $RUNNER_OS == "Linux" ]; then
-    CFLAGS="-fopenmp" LDFLAGS="-lgomp" uv pip install .
+    CFLAGS="-fopenmp" LDFLAGS="-lgomp" uv pip install -e .
     export OMP_NUM_THREADS=4
 else
     uv pip list
-    uv pip install .
+    uv pip install -e .
 fi
 
 python -c "import cvxpy; print(cvxpy.installed_solvers())"

cvxpy/atoms/__init__.py

@@ -34,6 +34,7 @@
 from cvxpy.atoms.affine.reshape import deep_flatten, reshape
 from cvxpy.atoms.affine.squeeze import squeeze
 from cvxpy.atoms.affine.concatenate import concatenate
+from cvxpy.atoms.affine.stack import stack
 from cvxpy.atoms.affine.sum import sum
 from cvxpy.atoms.affine.trace import trace, Trace
 from cvxpy.atoms.affine.transpose import (transpose, permute_dims, swapaxes, moveaxis)

cvxpy/atoms/affine/stack.py

@@ -0,0 +1,107 @@
+"""
+Copyright, the CVXPY authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+NumPy-style stacking that inserts a new axis for CVXPY Expressions.
+
+This mirrors the semantics of ``np.stack``, where they make sense, for symbolic
+expressions: every input must share an identical shape, and the result gains a
+new axis of length ``len(arrays)`` at the requested position. Options that are
+specific to NumPy's ndarray type system (for example ``out`` or ``dtype``) are
+intentionally unsupported.
+
+Examples
+--------
+>>> import cvxpy as cp
+>>> a = cp.Parameter((3,))
+>>> b = cp.Parameter((3,))
+>>> y = cp.stack([a, b], axis=0)
+>>> y.shape
+(2, 3)
+>>> z = cp.stack([a, b], axis=-1)
+>>> z.shape
+(3, 2)
+"""
+
+from typing import Iterable, Sequence
+
+from cvxpy.atoms.affine.concatenate import concatenate
+from cvxpy.atoms.affine.reshape import reshape
+from cvxpy.expressions.expression import Expression
+
+
+def _as_expression(obj) -> Expression:
+    """Cast scalars/arrays to ``Constant`` while leaving Expressions intact."""
+    return obj if isinstance(obj, Expression) else Expression.cast_to_const(obj)
+
+
+def stack(arrays: Sequence[object] | Iterable[object], axis: int = 0) -> Expression:
+    """Join a sequence of expressions along a new axis.
+
+    Parameters
+    ----------
+    arrays
+        Sequence of expressions (or array-likes) that all have the same shape.
+    axis
+        Index of the new axis in the result. Values in ``[-(ndim + 1), ndim``
+        ``+ 1)`` are accepted, following ``numpy.stack``.
+
+    Returns
+    -------
+    Expression
+        Expression whose shape equals the common input shape with the new axis
+        inserted at ``axis`` and length ``len(arrays)`` along that axis.
+
+    Raises
+    ------
+    TypeError
+        If ``axis`` is not an integer.
+    ValueError
+        If ``arrays`` is empty, shapes differ, or ``axis`` is out of bounds.
+    """
+    xs = [_as_expression(arg) for arg in arrays]
+    if not xs:
+        raise ValueError("need at least one array to stack")
+
+    if not isinstance(axis, int):
+        raise TypeError(f"axis must be an int; received {type(axis).__name__}")
+
+    shapes = {expr.shape for expr in xs}
+    if len(shapes) != 1:
+        raise ValueError(
+            "all input arrays must have the same shape; got "
+            f"{sorted(shapes)}"
+        )
+
+    base_shape = xs[0].shape
+    result_ndim = len(base_shape) + 1
+    if not (-result_ndim <= axis < result_ndim):
+        raise ValueError(
+            f"axis {axis} is out of bounds for result ndim {result_ndim}"
+        )
+
+    axis_index = axis if axis >= 0 else axis + result_ndim
+    # Slice the shape so we can splice a singleton axis where ``axis`` points.
+    prefix = base_shape[:axis_index]
+    suffix = base_shape[axis_index:]
+    # Reshape each argument to inject the new length-1 axis before concatenating.
+    reshaped = [
+        reshape(expr, prefix + (1,) + suffix, order='F')
+        for expr in xs
+    ]
+
+    return concatenate(reshaped, axis=axis_index)
+
+
+__all__ = ["stack"]

cvxpy/tests/atoms/affine/test_stack.py

@@ -0,0 +1,110 @@
+"""
+Copyright, the CVXPY authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Unit tests for the NumPy-style ``cp.stack`` helper. Each test focuses on a
+single aspect of the API contract so failures clearly indicate which axis
+handling or validation rule regressed.
+"""
+
+import numpy as np
+import pytest
+
+import cvxpy as cp
+
+
+def test_stack_1d_axis0() -> None:
+    """Two 1-D Parameters stack along axis 0 to form a 2x4 array."""
+    a = cp.Parameter((4,))
+    b = cp.Parameter((4,))
+    y = cp.stack([a, b], axis=0)
+    assert y.shape == (2, 4)
+
+
+def test_stack_1d_axis_last_numeric_parity() -> None:
+    """Scalar arrays stacked along the last axis match NumPy numerically."""
+    a = np.array([1.0, 2.0, 3.0])
+    b = np.array([4.0, 5.0, 6.0])
+    y = cp.stack([a, b], axis=-1)
+    assert y.shape == (3, 2)
+    expected = np.stack([a, b], axis=-1)
+    assert np.allclose(y.value, expected)
+
+
+def test_stack_2d_various_axes() -> None:
+    """Validate axis normalization on 2-D operands for several positions."""
+    a = cp.Parameter((3, 4))
+    b = cp.Parameter((3, 4))
+    assert cp.stack([a, b], axis=0).shape == (2, 3, 4)
+    assert cp.stack([a, b], axis=1).shape == (3, 2, 4)
+    assert cp.stack([a, b], axis=-1).shape == (3, 4, 2)
+
+
+def test_stack_scalar_inputs() -> None:
+    """Literal scalars auto-wrap into Constants and stack as a 1-D vector."""
+    y = cp.stack([1, 2, 3], axis=0)
+    assert y.shape == (3,)
+    assert np.allclose(y.value, np.array([1, 2, 3]))
+
+
+def test_stack_shape_mismatch_raises() -> None:
+    """Inputs of different shapes trigger the NumPy-style ValueError."""
+    a = cp.Parameter((3,))
+    b = cp.Parameter((4,))
+    with pytest.raises(ValueError):
+        cp.stack([a, b], axis=0)
+
+
+def test_stack_empty_list_raises() -> None:
+    """An empty input sequence is rejected."""
+    with pytest.raises(ValueError):
+        cp.stack([], axis=0)
+
+
+def test_stack_axis_bounds_check() -> None:
+    """Axis validation mirrors NumPy bounds for the resulting ndim."""
+    a = cp.Parameter((3,))
+    # ndim=1 -> result_ndim=2; valid axes: -2, -1, 0, 1
+    for good in (-2, -1, 0, 1):
+        cp.stack([a, a], axis=good)
+    for bad in (-3, 2, 3):
+        with pytest.raises(ValueError):
+            cp.stack([a, a], axis=bad)
+
+
+def test_stack_non_int_axis_raises() -> None:
+    """Non-integer axes raise a TypeError before shape checks run."""
+    a = cp.Parameter((2,))
+    with pytest.raises(TypeError):
+        cp.stack([a, a], axis=0.5)
+
+
+def test_stack_variables_shape_only() -> None:
+    """Variables share shape metadata with the stacked expression."""
+    x = cp.Variable((2, 3))
+    y = cp.Variable((2, 3))
+    z = cp.stack([x, y], axis=0)
+    assert z.shape == (2, 2, 3)
+
+
+def test_stack_canonicalization_resolves_equalities() -> None:
+    """Canonicalization maps scalar variables onto the stacked vector."""
+    x = cp.Variable()
+    y = cp.Variable()
+    z = cp.Variable(2)
+    z_tilde = cp.stack([x, y])
+    problem = cp.Problem(cp.Minimize(0), [z_tilde == z, x == 1, y == 2])
+    problem.solve(solver=cp.SCS)
+    assert problem.status == cp.OPTIMAL
+    assert np.allclose(z.value, np.array([1.0, 2.0]))

doc/source/api_reference/cvxpy.atoms.affine.rst

@@ -271,4 +271,11 @@ vec_to_upper_tri
 vstack
 -----------------------------------
 
-.. autofunction:: cvxpy.vstack
+.. autofunction:: cvxpy.vstack
+
+.. _stack:
+
+stack
+-----------------------------------
+
+.. autofunction:: cvxpy.stack