cp.sum support generator (cvxpy#3058)

* document current sum of list behaviour

* support generator input

* use GeneratorType to make generator support cleaner

---------

Co-authored-by: Muhammad Yasirroni <yasirroni@users.noreply.github.com>

Author: yasirroni

cvxpy/atoms/affine/sum.py

@@ -15,6 +15,7 @@
 """
 import builtins
 from functools import wraps
+from types import GeneratorType
 from typing import Optional, Tuple
 
 import numpy as np
@@ -145,7 +146,7 @@ def sum(expr, axis: Optional[int] = None, keepdims: bool = False):
     """
     Wrapper for Sum class.
     """
-    if isinstance(expr, list):
+    if isinstance(expr, (list, GeneratorType)):
         return builtins.sum(expr)
     else:
         return Sum(expr, axis, keepdims)

cvxpy/tests/test_atoms.py

@@ -492,6 +492,14 @@ def test_sum(self) -> None:
         self.assertEqual(cp.sum(Variable(2)).shape, tuple())
         self.assertEqual(cp.sum(Variable(2)).curvature, s.AFFINE)
         self.assertEqual(cp.sum(Variable((2, 1)), keepdims=True).shape, (1, 1))
+
+        # Iterables and Generators
+        self.assertEqual(cp.sum([Variable(1) for _ in range(3)]).shape, (1,))
+        self.assertEqual(cp.sum([Variable(2) for _ in range(3)]).shape, (2,))
+        self.assertEqual(cp.sum(Variable(1) for _ in range(3)).shape, (1,))
+        self.assertEqual(cp.sum(Variable(2) for _ in range(3)).shape, (2,))
+        self.assertEqual(cp.sum(range(3)).shape, tuple())
+
         # Mixed curvature.
         mat = np.array([[1, -1]])
         self.assertEqual(cp.sum(mat @ cp.square(Variable(2))).curvature, s.UNKNOWN)