Adding `einsum` atom #2970

NDevanathan · 2025-10-14T00:53:29Z

Description

I implemented an einsum atom (based on numpy.einsum). The atom is affine and provides a simple interface for describing linear operations on tensors in CVXPY expressions. I wrote tests for the atom in cvxpy/tests/test_einsum.py. All but one of the tests is passing.

The last block of asserts in test_einsum_solve() in cvxpy/tests/test_einsum.py fails (at solve time) non-deterministically on my machine (it seems to fail roughly 50% of the time). Specifically, I get an error reshaping a CVXPY expression with shape (4,) to shape (4,1), and I am unsure of how to resolve this issue.

Issue link (if applicable): N/A

Type of change

New feature (backwards compatible)
New feature (breaking API changes)
Bug fix
Other (Documentation, CI, ...)

Contribution checklist

Add our license to new files.
Check that your code adheres to our coding style.
Write unittests.
Run the unittests and check that they’re passing.
Run the benchmarks to make sure your change doesn’t introduce a regression.

CLAassistant · 2025-10-14T00:53:35Z

All committers have signed the CLA.

github-actions · 2025-10-14T18:11:05Z

Benchmarks that have stayed the same:

   before           after         ratio
 [c78d671a]       [123d7f69]
      1.37±0s          1.46±0s     1.07  matrix_stuffing.ParamConeMatrixStuffing.time_compile_problem
      558±0ms          588±0ms     1.05  simple_QP_benchmarks.ParametrizedQPBenchmark.time_compile_problem
      278±0ms          292±0ms     1.05  matrix_stuffing.ParamSmallMatrixStuffing.time_compile_problem
      1.08±0s          1.12±0s     1.03  gini_portfolio.Cajas.time_compile_problem
      2.32±0s          2.38±0s     1.03  simple_LP_benchmarks.SimpleFullyParametrizedLPBenchmark.time_compile_problem
      1.65±0s          1.69±0s     1.02  tv_inpainting.TvInpainting.time_compile_problem
      12.9±0s          13.1±0s     1.02  finance.CVaRBenchmark.time_compile_problem
      520±0ms          527±0ms     1.01  semidefinite_programming.SemidefiniteProgramming.time_compile_problem
      233±0ms          236±0ms     1.01  gini_portfolio.Murray.time_compile_problem
     38.7±0ms         39.1±0ms     1.01  matrix_stuffing.SmallMatrixStuffing.time_compile_problem
      22.4±0s          22.6±0s     1.01  sdp_segfault_1132_benchmark.SDPSegfault1132Benchmark.time_compile_problem
      1.09±0s          1.10±0s     1.00  finance.FactorCovarianceModel.time_compile_problem
      340±0ms          340±0ms     1.00  gini_portfolio.Yitzhaki.time_compile_problem
      4.95±0s          4.95±0s     1.00  optimal_advertising.OptimalAdvertising.time_compile_problem
      243±0ms          243±0ms     1.00  high_dim_convex_plasticity.ConvexPlasticity.time_compile_problem
      11.4±0s          11.3±0s     1.00  simple_LP_benchmarks.SimpleLPBenchmark.time_compile_problem
      294±0ms          292±0ms     1.00  slow_pruning_1668_benchmark.SlowPruningBenchmark.time_compile_problem
      4.64±0s          4.62±0s     0.99  huber_regression.HuberRegression.time_compile_problem
      256±0ms          255±0ms     0.99  simple_QP_benchmarks.SimpleQPBenchmark.time_compile_problem
      5.17±0s          5.12±0s     0.99  svm_l1_regularization.SVMWithL1Regularization.time_compile_problem
      827±0ms          819±0ms     0.99  simple_QP_benchmarks.LeastSquares.time_compile_problem
      677±0ms          669±0ms     0.99  matrix_stuffing.ConeMatrixStuffingBench.time_compile_problem
      2.87±0s          2.83±0s     0.99  quantum_hilbert_matrix.QuantumHilbertMatrix.time_compile_problem
      1.88±0s          1.86±0s     0.99  simple_QP_benchmarks.UnconstrainedQP.time_compile_problem
      952±0ms          937±0ms     0.98  simple_LP_benchmarks.SimpleScalarParametrizedLPBenchmark.time_compile_problem

Transurgeon

@NDevanathan sorry it took so long to review your code. I noticed one small nitpicky import in your einsum implementation.
Otherwise I think overall the code is a bit difficult to follow but definitely seems to be correct (seeing all the tests you wrote pass). We still need to understand why the flaky test is happening (it is really so weird! could it be because we are importing some numpy core functions?).
One thing on my end would be to add more docstrings and also add the einsum atom to the documentation somewhere (there is a table and also some API references to update).

Thanks for doing this, really great work!

Transurgeon · 2025-10-20T15:05:28Z

cvxpy/atoms/affine/einsum.py

+
+from cvxpy.atoms.affine.binary_operators import (
+    multiply,
+    reshape,


I just checked and there is no reshape in binary_operators? Are you sure you are importing the right one?

Good catch. I fixed the reshape import.

NDevanathan · 2025-10-20T16:17:58Z

One thing on my end would be to add more docstrings and also add the einsum atom to the documentation somewhere (there is a table and also some API references to update).

I updated the docstrings to be better. I will try to add einsum to the docs.

NDevanathan · 2025-10-20T16:19:43Z

Otherwise I think overall the code is a bit difficult to follow but definitely seems to be correct (seeing all the tests you wrote pass). We still need to understand why the flaky test is happening (it is really so weird! could it be because we are importing some numpy core functions?).

I think there might be an issue with how I implement einsum and how DGP problems are canonicalized. The error occurs at solve time, and DCP solves work. So, I think the issue occurs in doing the DGP canonicalization chain of einsum. I will also try to dig into this issue.

NDevanathan · 2025-12-21T02:47:15Z

Sorry for the great delay, I have fixed the issues that were leading to test failures I think @Transurgeon . I will add docs today or tomorrow.

cvxpy/atoms/affine/einsum.py

PTNobel

Awesome work Nikhil; the tests look great.

Here's what I'd like to see:

Some plan to handle if NumPy changes their internal function names in the future (see below)
A tad more documentation at the top of einsum on either the algorithm being applied, or just linking to a reference.

cvxpy/atoms/affine/einsum.py

PTNobel · 2025-12-22T00:56:21Z

cvxpy/atoms/affine/einsum.py

+    Expression
+        The contracted expression.
+    """
+    # Initial parsing


Can you add a comment that provides the overview of what happens in this function? It can be a reference to external documentation or what the implementation is based on, but this is pretty opaque and involves a lot of complicated functions.

Addressed in most recent commit.

…e thoroughly

PTNobel

Last two comments!

PTNobel · 2025-12-22T03:12:47Z

cvxpy/atoms/affine/einsum.py

    is cubic in the number of distinct subscripts. We typically expect the greedy 
    search to produce the optimal path for most problems.

+    Here, einsum is implemented using the CVXPY sum, multiply, permute_dims, and reshape atoms.


This doesn't need to be in the doc comment! This is just for devs, so you can put it in a comment block below the doc comment!

cvxpy/atoms/affine/einsum.py

PTNobel · 2025-12-22T03:19:50Z

cvxpy/atoms/affine/einsum.py

+        input_subscripts, output_subscript = subscripts.split("->")
+        ```
+        The rest is validation.
+    


There are empty spaces on this line.

PTNobel · 2025-12-22T03:20:00Z

cvxpy/atoms/affine/einsum.py

+    """Note for maintainers:
+    Here, einsum is implemented using the CVXPY sum, multiply, permute_dims, and reshape atoms.
+    The implementation proceeds as follows:
+    


NDevanathan added 2 commits October 14, 2025 00:43

Adding einsum atom; squashed commits

6e199c6

Merge branch 'cvxpy:master' into master

57c6664

Transurgeon added the PR no backport needed label Oct 14, 2025

isolated DGP test

9fab1e9

Transurgeon reviewed Oct 20, 2025

View reviewed changes

removed incorrect import, updated docstrings

806db40

NDevanathan added 2 commits December 20, 2025 17:21

Merge branch 'cvxpy:master' into master

3384637

fixing tricky canonicalization issue

0242afd

adding docs

29ace05

PTNobel reviewed Dec 21, 2025

View reviewed changes

cvxpy/atoms/affine/einsum.py Outdated Show resolved Hide resolved

PTNobel requested changes Dec 22, 2025

View reviewed changes

removing dependance on numpy internals; documented implementation mor…

d3cd3e1

…e thoroughly

PTNobel approved these changes Dec 22, 2025

View reviewed changes

minor edit to einsum docstring

ee11455

PTNobel reviewed Dec 22, 2025

View reviewed changes

ran ruff format

a4898d9

PTNobel merged commit 0822817 into cvxpy:master Dec 22, 2025
42 of 44 checks passed

Adding einsum atom #2970

Adding einsum atom #2970

Uh oh!

Conversation

NDevanathan commented Oct 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Contribution checklist

Uh oh!

CLAassistant commented Oct 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Oct 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Transurgeon left a comment

Choose a reason for hiding this comment

Uh oh!

Transurgeon Oct 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NDevanathan Oct 20, 2025

Choose a reason for hiding this comment

Uh oh!

NDevanathan commented Oct 20, 2025

Uh oh!

NDevanathan commented Oct 20, 2025

Uh oh!

NDevanathan commented Dec 21, 2025

Uh oh!

Uh oh!

PTNobel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

PTNobel Dec 22, 2025

Choose a reason for hiding this comment

Uh oh!

NDevanathan Dec 22, 2025

Choose a reason for hiding this comment

Uh oh!

PTNobel left a comment

Choose a reason for hiding this comment

Uh oh!

PTNobel Dec 22, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

PTNobel Dec 22, 2025

Choose a reason for hiding this comment

Uh oh!

PTNobel Dec 22, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Adding `einsum` atom #2970

Adding `einsum` atom #2970

NDevanathan commented Oct 14, 2025 •

edited

Loading

CLAassistant commented Oct 14, 2025 •

edited

Loading

github-actions bot commented Oct 14, 2025 •

edited

Loading

Transurgeon Oct 20, 2025 •

edited

Loading