ENH: einsum: Disable broadcasting by default, allow spaces in subscripts string

3 files changed, 138 insertions, 81 deletions

diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index 51826c5ff..f51860240 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -1540,13 +1540,16 @@ add_newdoc('numpy.core', 'einsum',
     ``np.einsum('ji', a)`` takes its transpose.
 
     The output can be controlled by specifying output subscript labels
-    as well.  This specifies the label order, and allows summing to be
-    disallowed or forced when desired.  The call ``np.einsum('i->', a)``
-    is equivalent to ``np.sum(a, axis=-1)``, and
-    ``np.einsum('ii->i', a)`` is equivalent to ``np.diag(a)``.
-
-    It is also possible to control how broadcasting occurs using
-    an ellipsis.  To take the trace along the first and last axes,
+    as well.  This specifies the label order, and allows summing to
+    be disallowed or forced when desired.  The call ``np.einsum('i->', a)``
+    is like ``np.sum(a, axis=-1)``, and ``np.einsum('ii->i', a)``
+    is like ``np.diag(a)``.  The difference is that ``einsum`` does not
+    allow broadcasting by default.
+
+    To enable and control broadcasting, use an ellipsis.  Default
+    NumPy-style broadcasting is done by adding an ellipsis
+    to the left of each term, like ``np.einsum('...ii->...i', a)``.
+    To take the trace along the first and last axes,
     you can do ``np.einsum('i...i', a)``, or to do a matrix-matrix
     product with the left-most indices instead of rightmost, you can do
     ``np.einsum('ij...,jk...->ik...', a, b)``.
@@ -1624,7 +1627,7 @@ add_newdoc('numpy.core', 'einsum',
            [1, 4],
            [2, 5]])
 
-    >>> np.einsum(',', 3, c)
+    >>> np.einsum('..., ...', 3, c)
     array([[ 0,  3,  6],
            [ 9, 12, 15]])
     >>> np.multiply(3, c)
@@ -1643,7 +1646,7 @@ add_newdoc('numpy.core', 'einsum',
     array([[0, 1, 2, 3, 4],
            [0, 2, 4, 6, 8]])
 
-    >>> np.einsum('i...->', a)
+    >>> np.einsum('i...->...', a)
     array([50, 55, 60, 65, 70])
     >>> np.sum(a, axis=0)
     array([50, 55, 60, 65, 70])
diff --git a/numpy/core/src/multiarray/einsum.c.src b/numpy/core/src/multiarray/einsum.c.src
index f5c91b170..1c3b07bde 100644
--- a/numpy/core/src/multiarray/einsum.c.src
+++ b/numpy/core/src/multiarray/einsum.c.src
@@ -46,6 +46,7 @@
 #define EINSUM_IS_SSE_ALIGNED(x) ((((npy_intp)x)&0xf) == 0)
 
 typedef enum {
+    BROADCAST_NONE,
     BROADCAST_LEFT,
     BROADCAST_RIGHT,
     BROADCAST_MIDDLE
@@ -1382,7 +1383,7 @@ parse_operand_subscripts(char *subscripts, int length,
                         EINSUM_BROADCAST *out_broadcast)
 {
     int i, idim, ndim_left, label;
-    int left_labels = 0, right_labels = 0;
+    int left_labels = 0, right_labels = 0, ellipsis = 0;
 
     /* Process the labels from the end until the ellipsis */
     idim = ndim-1;
@@ -1417,6 +1418,7 @@ parse_operand_subscripts(char *subscripts, int length,
         else if (label == '.') {
             /* A valid ellipsis */
             if (i >= 2 && subscripts[i-1] == '.' && subscripts[i-2] == '.') {
+                ellipsis = 1;
                 length = i-2;
                 break;
             }
@@ -1428,7 +1430,7 @@ parse_operand_subscripts(char *subscripts, int length,
                 
             }
         }
-        else {
+        else if (label != ' ') {
             PyErr_Format(PyExc_ValueError,
                         "invalid subscript '%c' in einstein sum "
                         "subscripts string, subscripts must "
@@ -1436,6 +1438,15 @@ parse_operand_subscripts(char *subscripts, int length,
             return 0;
         }
     }
+
+    if (!ellipsis && idim != -1) {
+        PyErr_Format(PyExc_ValueError,
+                    "operand has more dimensions than subscripts "
+                    "given in einstein sum, but no '...' ellipsis "
+                    "provided to broadcast the extra dimensions.");
+        return 0;
+    }
+
     /* Reduce ndim to just the dimensions left to fill at the beginning */
     ndim_left = idim+1;
     idim = 0;
@@ -1472,7 +1483,7 @@ parse_operand_subscripts(char *subscripts, int length,
                     return 0;
                 }
             }
-            else {
+            else if (label != ' ') {
                 PyErr_Format(PyExc_ValueError,
                             "invalid subscript '%c' in einstein sum "
                             "subscripts string, subscripts must "
@@ -1509,7 +1520,10 @@ parse_operand_subscripts(char *subscripts, int length,
         }
     }
 
-    if (left_labels && right_labels) {
+    if (!ellipsis) {
+        *out_broadcast = BROADCAST_NONE;
+    }
+    else if (left_labels && right_labels) {
         *out_broadcast = BROADCAST_MIDDLE;
     }
     else if (!left_labels) {
@@ -1535,7 +1549,7 @@ parse_output_subscripts(char *subscripts, int length,
                         EINSUM_BROADCAST *out_broadcast)
 {
     int i, nlabels, label, idim, ndim, ndim_left;
-    int left_labels = 0, right_labels = 0;
+    int left_labels = 0, right_labels = 0, ellipsis = 0;
 
     /* Count the labels, making sure they're all unique and valid */
     nlabels = 0;
@@ -1563,7 +1577,7 @@ parse_output_subscripts(char *subscripts, int length,
                 return -1;
             }
         }
-        else if (label != '.') {
+        else if (label != '.' && label != ' ') {
             PyErr_Format(PyExc_ValueError,
                         "invalid subscript '%c' in einstein sum "
                         "subscripts string, subscripts must "
@@ -1580,7 +1594,7 @@ parse_output_subscripts(char *subscripts, int length,
     for (i = length-1; i >= 0; --i) {
         label = subscripts[i];
         /* A label for an axis */
-        if (label != '.') {
+        if (label != '.' && label != ' ') {
             if (idim >= 0) {
                 out_labels[idim--] = label;
             }
@@ -1593,9 +1607,10 @@ parse_output_subscripts(char *subscripts, int length,
             right_labels = 1;
         }
         /* The end of the ellipsis */
-        else {
+        else if (label == '.') {
             /* A valid ellipsis */
             if (i >= 2 && subscripts[i-1] == '.' && subscripts[i-2] == '.') {
+                ellipsis = 1;
                 length = i-2;
                 break;
             }
@@ -1608,6 +1623,15 @@ parse_output_subscripts(char *subscripts, int length,
             }
         }
     }
+
+    if (!ellipsis && idim != -1) {
+        PyErr_SetString(PyExc_ValueError,
+                    "output has more dimensions than subscripts "
+                    "given in einstein sum, but no '...' ellipsis "
+                    "provided to broadcast the extra dimensions.");
+        return 0;
+    }
+
     /* Reduce ndim to just the dimensions left to fill at the beginning */
     ndim_left = idim+1;
     idim = 0;
@@ -1620,7 +1644,7 @@ parse_output_subscripts(char *subscripts, int length,
         for (i = 0; i < length; ++i) {
             label = subscripts[i];
             /* A label for an axis */
-            if (label != '.') {
+            if (label != '.' && label != ' ') {
                 if (idim < ndim_left) {
                     out_labels[idim++] = label;
                 }
@@ -1646,7 +1670,10 @@ parse_output_subscripts(char *subscripts, int length,
         out_labels[idim++] = 0;
     }
 
-    if (left_labels && right_labels) {
+    if (!ellipsis) {
+        *out_broadcast = BROADCAST_NONE;
+    }
+    else if (left_labels && right_labels) {
         *out_broadcast = BROADCAST_MIDDLE;
     }
     else if (!left_labels) {
@@ -1941,7 +1968,7 @@ prepare_op_axes(int ndim, int iop, char *labels, npy_intp *axes,
             }
         }
     }
-    /* Middle broadcasting */
+    /* Middle or None broadcasting */
     else {
         /* broadcast dimensions get placed in leftmost position */
         ibroadcast = 0; 
@@ -2133,8 +2160,13 @@ PyArray_EinsteinSum(char *subscripts, npy_intp nop,
      * that appeared once, in alphabetical order
      */
     if (subscripts[0] == '\0') {
-        char outsubscripts[NPY_MAXDIMS];
-        int length = 0;
+        char outsubscripts[NPY_MAXDIMS + 3];
+        int length;
+        /* If no output was specified, always broadcast left (like normal) */
+        outsubscripts[0] = '.';
+        outsubscripts[1] = '.';
+        outsubscripts[2] = '.';
+        length = 3;
         for (label = min_label; label <= max_label; ++label) {
             if (label_counts[label] == 1) {
                 if (length < NPY_MAXDIMS-1) {
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
index f1cf7c2c4..68b40ff9b 100644
--- a/numpy/core/tests/test_numeric.py
+++ b/numpy/core/tests/test_numeric.py
@@ -211,11 +211,16 @@ class TestEinSum(TestCase):
         assert_raises(ValueError, np.einsum, "ii", np.arange(6).reshape(2,3))
         assert_raises(ValueError, np.einsum, "ii->i", np.arange(6).reshape(2,3))
 
+        # broadcasting to new dimensions must be enabled explicitly
+        assert_raises(ValueError, np.einsum, "i", np.arange(6).reshape(2,3))
+        assert_raises(ValueError, np.einsum, "i->i", [[0,1],[0,1]],
+                                            out=np.arange(4).reshape(2,2))
+
     def test_einsum_views(self):
         # pass-through
         a = np.arange(6).reshape(2,3)
 
-        b = np.einsum("", a)
+        b = np.einsum("...", a)
         assert_(b.base is a)
 
         b = np.einsum("ij", a)
@@ -239,16 +244,16 @@ class TestEinSum(TestCase):
         # diagonal with various ways of broadcasting an additional dimension
         a = np.arange(27).reshape(3,3,3)
 
-        b = np.einsum("ii->i", a)
+        b = np.einsum("...ii->...i", a)
         assert_(b.base is a)
         assert_equal(b, [[x[i,i] for i in range(3)] for x in a])
 
-        b = np.einsum("ii...->i", a)
+        b = np.einsum("ii...->...i", a)
         assert_(b.base is a)
         assert_equal(b, [[x[i,i] for i in range(3)]
                          for x in a.transpose(2,0,1)])
 
-        b = np.einsum("ii->i...", a)
+        b = np.einsum("...ii->i...", a)
         assert_(b.base is a)
         assert_equal(b, [a[:,i,i] for i in range(3)])
 
@@ -264,7 +269,7 @@ class TestEinSum(TestCase):
         assert_(b.base is a)
         assert_equal(b, [a.transpose(1,0,2)[:,i,i] for i in range(3)])
 
-        b = np.einsum("i...i->i", a)
+        b = np.einsum("i...i->...i", a)
         assert_(b.base is a)
         assert_equal(b, [[x[i,i] for i in range(3)]
                          for x in a.transpose(1,0,2)])
@@ -288,33 +293,42 @@ class TestEinSum(TestCase):
         a = np.arange(10, dtype=dtype)
         assert_equal(np.einsum("i->", a), np.sum(a, axis=-1))
 
-        a = np.arange(24, dtype=dtype).reshape(2,3,4)
-        assert_equal(np.einsum("i->", a), np.sum(a, axis=-1))
+        for n in range(1,17):
+            a = np.arange(2*3*n, dtype=dtype).reshape(2,3,n)
+            assert_equal(np.einsum("...i->...", a),
+                         np.sum(a, axis=-1).astype(dtype))
 
         # sum(a, axis=0)
-        a = np.arange(10, dtype=dtype)
-        assert_equal(np.einsum("i...->", a), np.sum(a, axis=0))
+        for n in range(1,17):
+            a = np.arange(2*n, dtype=dtype).reshape(2,n)
+            assert_equal(np.einsum("i...->...", a),
+                         np.sum(a, axis=0).astype(dtype))
 
-        a = np.arange(24, dtype=dtype).reshape(2,3,4)
-        assert_equal(np.einsum("i...->", a), np.sum(a, axis=0))
+        for n in range(1,17):
+            a = np.arange(2*3*n, dtype=dtype).reshape(2,3,n)
+            assert_equal(np.einsum("i...->...", a),
+                         np.sum(a, axis=0).astype(dtype))
 
         # trace(a)
         a = np.arange(25, dtype=dtype).reshape(5,5)
         assert_equal(np.einsum("ii", a), np.trace(a))
 
         # multiply(a, b)
-        a = np.arange(12, dtype=dtype).reshape(3,4)
-        b = np.arange(24, dtype=dtype).reshape(2,3,4)
-        assert_equal(np.einsum(",", a, b), np.multiply(a, b))
+        for n in range(1,17):
+            a = np.arange(3*n, dtype=dtype).reshape(3,n)
+            b = np.arange(2*3*n, dtype=dtype).reshape(2,3,n)
+            assert_equal(np.einsum("..., ...", a, b), np.multiply(a, b))
 
         # inner(a,b)
-        a = np.arange(24, dtype=dtype).reshape(2,3,4)
-        b = np.arange(4, dtype=dtype)
-        assert_equal(np.einsum("i,i", a, b), np.inner(a, b))
+        for n in range(1,17):
+            a = np.arange(2*3*n, dtype=dtype).reshape(2,3,n)
+            b = np.arange(n, dtype=dtype)
+            assert_equal(np.einsum("...i, ...i", a, b), np.inner(a, b))
 
-        a = np.arange(24, dtype=dtype).reshape(2,3,4)
-        b = np.arange(2, dtype=dtype)
-        assert_equal(np.einsum("i...,i...", a, b), np.inner(a.T, b.T).T)
+        for n in range(1,11):
+            a = np.arange(n*3*2, dtype=dtype).reshape(n,3,2)
+            b = np.arange(n, dtype=dtype)
+            assert_equal(np.einsum("i..., i...", a, b), np.inner(a.T, b.T).T)
 
         # outer(a,b)
         a = np.arange(3, dtype=dtype)+1
@@ -328,28 +342,33 @@ class TestEinSum(TestCase):
             warnings.simplefilter('ignore', np.ComplexWarning)
 
             # matvec(a,b) / a.dot(b) where a is matrix, b is vector
-            a = np.arange(20, dtype=dtype).reshape(4,5)
-            b = np.arange(5, dtype=dtype)
-            assert_equal(np.einsum("ij,j", a, b), np.dot(a, b))
-
-            a = np.arange(20, dtype=dtype).reshape(4,5)
-            b = np.arange(5, dtype=dtype)
-            c = np.arange(4, dtype=dtype)
-            np.einsum("ij,j", a, b, out=c,
-                        dtype='f8', casting='unsafe')
-            assert_equal(c,
-                        np.dot(a.astype('f8'), b.astype('f8')).astype(dtype))
-
-            a = np.arange(20, dtype=dtype).reshape(4,5)
-            b = np.arange(5, dtype=dtype)
-            assert_equal(np.einsum("ji,j", a.T, b.T), np.dot(b.T, a.T))
-
-            a = np.arange(20, dtype=dtype).reshape(4,5)
-            b = np.arange(5, dtype=dtype)
-            c = np.arange(4, dtype=dtype)
-            np.einsum("ji,j", a.T, b.T, out=c, dtype='f8', casting='unsafe')
-            assert_equal(c,
-                    np.dot(b.T.astype('f8'), a.T.astype('f8')).astype(dtype))
+            for n in range(1,17):
+                a = np.arange(4*n, dtype=dtype).reshape(4,n)
+                b = np.arange(n, dtype=dtype)
+                assert_equal(np.einsum("ij, j", a, b), np.dot(a, b))
+
+            for n in range(1,17):
+                a = np.arange(4*n, dtype=dtype).reshape(4,n)
+                b = np.arange(n, dtype=dtype)
+                c = np.arange(4, dtype=dtype)
+                np.einsum("ij,j", a, b, out=c,
+                            dtype='f8', casting='unsafe')
+                assert_equal(c,
+                            np.dot(a.astype('f8'),
+                                   b.astype('f8')).astype(dtype))
+
+            for n in range(1,17):
+                a = np.arange(4*n, dtype=dtype).reshape(4,n)
+                b = np.arange(n, dtype=dtype)
+                assert_equal(np.einsum("ji,j", a.T, b.T), np.dot(b.T, a.T))
+
+                a = np.arange(4*n, dtype=dtype).reshape(4,n)
+                b = np.arange(n, dtype=dtype)
+                c = np.arange(4, dtype=dtype)
+                np.einsum("ji,j", a.T, b.T, out=c, dtype='f8', casting='unsafe')
+                assert_equal(c,
+                        np.dot(b.T.astype('f8'),
+                               a.T.astype('f8')).astype(dtype))
 
             # matmat(a,b) / a.dot(b) where a is matrix, b is matrix
             a = np.arange(20, dtype=dtype).reshape(4,5)
@@ -363,7 +382,7 @@ class TestEinSum(TestCase):
             assert_equal(c,
                         np.dot(a.astype('f8'), b.astype('f8')).astype(dtype))
 
-            # matrix triple product (note this is not an efficient
+            # matrix triple product (note this is not currently an efficient
             # way to multiply 3 matrices)
             a = np.arange(12, dtype=dtype).reshape(3,4)
             b = np.arange(20, dtype=dtype).reshape(4,5)
@@ -385,7 +404,7 @@ class TestEinSum(TestCase):
             if np.dtype(dtype) != np.dtype('f2'):
                 a = np.arange(60, dtype=dtype).reshape(3,4,5)
                 b = np.arange(24, dtype=dtype).reshape(4,3,2)
-                assert_equal(np.einsum("ijk,jil->kl", a, b),
+                assert_equal(np.einsum("ijk, jil -> kl", a, b),
                                 np.tensordot(a,b, axes=([1,0],[0,1])))
 
                 a = np.arange(60, dtype=dtype).reshape(3,4,5)
@@ -411,21 +430,24 @@ class TestEinSum(TestCase):
         assert_equal(np.einsum("i,->", a, 3), 3*np.sum(a))
 
         # Various stride0, contiguous, and SSE aligned variants
-        a = np.arange(64, dtype=dtype)
-        if np.dtype(dtype).itemsize > 1:
-            assert_equal(np.einsum(",",a,a), np.multiply(a,a))
-            assert_equal(np.einsum("i,i", a, a), np.dot(a,a))
-            assert_equal(np.einsum("i,->i", a, 2), 2*a)
-            assert_equal(np.einsum(",i->i", 2, a), 2*a)
-            assert_equal(np.einsum("i,->", a, 2), 2*np.sum(a))
-            assert_equal(np.einsum(",i->", 2, a), 2*np.sum(a))
-
-            assert_equal(np.einsum(",",a[1:],a[:-1]), np.multiply(a[1:],a[:-1]))
-            assert_equal(np.einsum("i,i", a[1:], a[:-1]), np.dot(a[1:],a[:-1]))
-            assert_equal(np.einsum("i,->i", a[1:], 2), 2*a[1:])
-            assert_equal(np.einsum(",i->i", 2, a[1:]), 2*a[1:])
-            assert_equal(np.einsum("i,->", a[1:], 2), 2*np.sum(a[1:]))
-            assert_equal(np.einsum(",i->", 2, a[1:]), 2*np.sum(a[1:]))
+        for n in range(1,25):
+            a = np.arange(n, dtype=dtype)
+            if np.dtype(dtype).itemsize > 1:
+                assert_equal(np.einsum("...,...",a,a), np.multiply(a,a))
+                assert_equal(np.einsum("i,i", a, a), np.dot(a,a))
+                assert_equal(np.einsum("i,->i", a, 2), 2*a)
+                assert_equal(np.einsum(",i->i", 2, a), 2*a)
+                assert_equal(np.einsum("i,->", a, 2), 2*np.sum(a))
+                assert_equal(np.einsum(",i->", 2, a), 2*np.sum(a))
+
+                assert_equal(np.einsum("...,...",a[1:],a[:-1]),
+                             np.multiply(a[1:],a[:-1]))
+                assert_equal(np.einsum("i,i", a[1:], a[:-1]),
+                             np.dot(a[1:],a[:-1]))
+                assert_equal(np.einsum("i,->i", a[1:], 2), 2*a[1:])
+                assert_equal(np.einsum(",i->i", 2, a[1:]), 2*a[1:])
+                assert_equal(np.einsum("i,->", a[1:], 2), 2*np.sum(a[1:]))
+                assert_equal(np.einsum(",i->", 2, a[1:]), 2*np.sum(a[1:]))
 
         # An object array, summed as the data type
         a = np.arange(9, dtype=object)