DOC: core: Document the new einsum function

4 files changed, 152 insertions, 2 deletions

diff --git a/doc/source/reference/routines.linalg.rst b/doc/source/reference/routines.linalg.rst
index 4c3c676d9..173a6ad53 100644
--- a/doc/source/reference/routines.linalg.rst
+++ b/doc/source/reference/routines.linalg.rst
@@ -15,6 +15,7 @@ Matrix and vector products
    inner
    outer
    tensordot
+   einsum
    linalg.matrix_power
    kron
 
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index db84046cb..854338a9c 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -826,6 +826,7 @@ add_newdoc('numpy.core', 'inner',
     --------
     tensordot : Sum products over arbitrary axes.
     dot : Generalised matrix product, using second last dimension of `b`.
+    einsum : Einstein summation convention.
 
     Notes
     -----
@@ -1429,6 +1430,7 @@ add_newdoc('numpy.core', 'dot',
     --------
     vdot : Complex-conjugating dot product.
     tensordot : Sum products over arbitrary axes.
+    einsum : Einstein summation convention.
 
     Examples
     --------
@@ -1457,6 +1459,149 @@ add_newdoc('numpy.core', 'dot',
 
     """)
 
+add_newdoc('numpy.core', 'einsum',
+    """
+    einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='safe')
+
+    Evaluates the Einstein summation convention on the operands.
+
+    Using the Einstein summation convention, many common multi-dimensional
+    array operations can be represented in a simple fashion.  This function
+    provides a way compute such summations.
+
+    The best way to understand this function is to try the examples below,
+    which show how many common NumPy functions can be implemented as
+    calls to einsum.
+    
+    The subscripts string is a comma-separated list of subscript labels,
+    where each label refers to a dimension of the corresponding operand.
+    Repeated subscripts labels in one operand take the diagonal.  For example,
+    ``np.einsum('ii', a)`` is equivalent to ``np.trace(a)``.
+    
+    Whenever a label is repeated, it is summed, so ``np.einsum('i,i', a, b)``
+    is equivalent to ``np.inner(a,b)``.  If a label appears only once,
+    it is not summed, so ``np.einsum('i', a)`` produces a view of ``a``
+    with no changes.
+
+    The order of labels in the output is by default alphabetical.  This
+    means that ``np.einsum('ij', a)`` doesn't affect a 2D array, while
+    ``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,
+    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)``.
+
+    When there is only one operand, no axes are summed, and no output
+    parameter is provided, a view into the operand is returned instead
+    of a new array.  Thus, taking the diagonal as ``np.einsum('ii->i', a)``
+    produces a view.
+
+    Parameters
+    ----------
+    subscripts : string
+        Specifies the subscripts for summation.
+    operands : list of array_like
+        These are the arrays for the operation.
+    out : None or array
+        If provided, the calculation is done into this array.
+    dtype : None or data type
+        If provided, forces the calculation to use the data type specified.
+    order : 'C', 'F', 'A', or 'K'
+        Controls the memory layout of the output.
+    casting : 'no', 'equiv', 'safe', 'same_kind', 'unsafe'
+        Controls what kind of data casting may occur.  Setting this to
+        'unsafe' is not recommended, as it can adversely affect accumulations.
+
+    Returns
+    -------
+    output : ndarray
+        The calculation based on the Einstein summation convention.
+
+    See Also
+    --------
+    dot, inner, outer, tensordot
+
+    
+    Examples
+    --------
+
+    >>> a = np.arange(25).reshape(5,5)
+    >>> b = np.arange(5)
+    >>> c = np.arange(6).reshape(2,3)
+
+    >>> np.einsum('ii', a)
+    60
+    >>> np.trace(a)
+    60
+
+    >>> np.einsum('ii->i', a)
+    array([ 0,  6, 12, 18, 24])
+    >>> np.diag(a)
+    array([ 0,  6, 12, 18, 24])
+
+    >>> np.einsum('ij,j', a, b)
+    array([ 30,  80, 130, 180, 230])
+    >>> np.dot(a, b)
+    array([ 30,  80, 130, 180, 230])
+
+    >>> np.einsum('ji', c)
+    array([[0, 3],
+           [1, 4],
+           [2, 5]])
+    >>> c.T
+    array([[0, 3],
+           [1, 4],
+           [2, 5]])
+
+    >>> np.einsum(',', 3, c)
+    array([[ 0,  3,  6],
+           [ 9, 12, 15]])
+    >>> np.multiply(3, c)
+    array([[ 0,  3,  6],
+           [ 9, 12, 15]])
+
+    >>> np.einsum('i,i', b, b)
+    30
+    >>> np.inner(b,b)
+    30
+
+    >>> np.einsum('i,j', np.arange(2)+1, b)
+    array([[0, 1, 2, 3, 4],
+           [0, 2, 4, 6, 8]])
+    >>> np.outer(np.arange(2)+1, b)
+    array([[0, 1, 2, 3, 4],
+           [0, 2, 4, 6, 8]])
+
+    >>> np.einsum('i...->', a)
+    array([50, 55, 60, 65, 70])
+    >>> np.sum(a, axis=0)
+    array([50, 55, 60, 65, 70])
+
+    >>> a = np.arange(60.).reshape(3,4,5)
+    >>> b = np.arange(24.).reshape(4,3,2)
+    >>> np.einsum('ijk,jil->kl', a, b)
+    array([[ 4400.,  4730.],
+           [ 4532.,  4874.],
+           [ 4664.,  5018.],
+           [ 4796.,  5162.],
+           [ 4928.,  5306.]])
+    >>> np.tensordot(a,b, axes=([1,0],[0,1]))
+    array([[ 4400.,  4730.],
+           [ 4532.,  4874.],
+           [ 4664.,  5018.],
+           [ 4796.,  5162.],
+           [ 4928.,  5306.]])
+
+    """)
+
 add_newdoc('numpy.core', 'alterdot',
     """
     Change `dot`, `vdot`, and `innerproduct` to use accelerated BLAS functions.
diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index 8cdb18435..43456b97f 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -817,6 +817,10 @@ def outer(a,b):
     out : ndarray, shape (M, N)
         ``out[i, j] = a[i] * b[j]``
 
+    See also
+    --------
+    numpy.inner, numpy.einsum
+
     References
     ----------
     .. [1] : G. H. Golub and C. F. van Loan, *Matrix Computations*, 3rd
@@ -907,7 +911,7 @@ def tensordot(a, b, axes=2):
 
     See Also
     --------
-    numpy.dot
+    numpy.dot, numpy.einsum
 
     Notes
     -----
diff --git a/numpy/linalg/linalg.py b/numpy/linalg/linalg.py
index c044176cf..c23338cd8 100644
--- a/numpy/linalg/linalg.py
+++ b/numpy/linalg/linalg.py
@@ -205,7 +205,7 @@ def tensorsolve(a, b, axes=None):
 
     See Also
     --------
-    tensordot, tensorinv
+    tensordot, tensorinv, einsum
 
     Examples
     --------