numpy / numpy/numpy

ENH: Adding reshape notation to einsum string

Open
#25,698 9 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
32.8k
Forks
12.8k
Avg merge
1d 7h
Merged PRs (30d)
197

Description

Proposed new feature or change:

I have implemented in about 30 LOC the following functionality for einsum:

import numpy as np
ab = np.random.normal(size=(2,3))
abcdefgh = np.random.normal(size=(2,3,4,5,6,7,8,9))
contracted = np.einsum('ab,abcdefgh->(cd)e(fgh)', ab, abcdefgh)  # <--- look at return string

assert contracted.shape == (4*5, 6, 7*8*9)
assert np.allclose(contracted, np.einsum('ab,abcdefgh->cdefgh', ab, abcdefgh).reshape(contracted.shape))

Note the return part of the string (cd)e(fgh): the idea is to use parentheses to enclose indices that we would like to reshape together.

I don't have time to follow up on docstring and tests. You can find the updated einsum function here

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the proposed implementation in numpy/_core/einsumfunc.py around line 1420 and review how existing einsum output notation is handled. Use the issue’s shape and np.allclose example to define expected behavior, then add the missing tests and docstring coverage noted by the author. Done means the parenthesized reshape notation works without regressing existing einsum behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.