Merge pull request #5878 from charris/quick-and-dirty-matmul

Quick and dirty matmul

9 files changed, 651 insertions, 29 deletions

diff --git a/doc/release/1.10.0-notes.rst b/doc/release/1.10.0-notes.rst
index f2ab6e99a..cb78b4e71 100644
--- a/doc/release/1.10.0-notes.rst
+++ b/doc/release/1.10.0-notes.rst
@@ -18,6 +18,7 @@ Highlights
   sequence of arrays along a new axis, complementing `np.concatenate` for
   joining along an existing axis.
 * Addition of `nanprod` to the set of nanfunctions.
+* Support for the '@' operator in Python 3.5.
 
 
 Dropped Support
@@ -153,6 +154,15 @@ vectors. An array of ``fweights`` indicates the number of repeats of each
 observation vector, and an array of ``aweights`` provides their relative
 importance or probability.
 
+Support for the '@' operator in Python 3.5+
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Python 3.5 adds support for a matrix multiplication operator '@' proposed
+in PEP465. Preliminary support for that has been implemented, and an
+equivalent function ``matmul`` has also been added for testing purposes and
+use in earlier Python versions. The function is preliminary and the order
+and number of its optional arguments can be expected to change.
+
+
 Improvements
 ============
 
diff --git a/doc/source/reference/arrays.ndarray.rst b/doc/source/reference/arrays.ndarray.rst
index c8d834d1c..e84b7fdf8 100644
--- a/doc/source/reference/arrays.ndarray.rst
+++ b/doc/source/reference/arrays.ndarray.rst
@@ -428,10 +428,10 @@ be performed.
    ndarray.all
    ndarray.any
 
-Arithmetic and comparison operations
-====================================
+Arithmetic, matrix multiplication, and comparison operations
+============================================================
 
-.. index:: comparison, arithmetic, operation, operator
+.. index:: comparison, arithmetic, matrix, operation, operator
 
 Arithmetic and comparison operations on :class:`ndarrays <ndarray>`
 are defined as element-wise operations, and generally yield
@@ -551,6 +551,20 @@ Arithmetic, in-place:
    casts the result to fit back in ``a``, whereas ``a = a + 3j``
    re-binds the name ``a`` to the result.
 
+Matrix Multiplication:
+
+.. autosummary::
+   :toctree: generated/
+
+   ndarray.__matmul__
+
+.. note::
+
+   Matrix operators ``@`` and ``@=`` were introduced in Python 3.5
+   following PEP465. Numpy 1.10 has a preliminary implementation of ``@``
+   for testing purposes. Further documentation can be found in the
+   :func:`matmul` documentation.
+
 
 Special methods
 ===============
diff --git a/doc/source/reference/routines.linalg.rst b/doc/source/reference/routines.linalg.rst
index 94533aaa9..bb2ad90a2 100644
--- a/doc/source/reference/routines.linalg.rst
+++ b/doc/source/reference/routines.linalg.rst
@@ -14,6 +14,7 @@ Matrix and vector products
    vdot
    inner
    outer
+   matmul
    tensordot
    einsum
    linalg.matrix_power
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index b0306957e..11a2688e5 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -1943,6 +1943,7 @@ add_newdoc('numpy.core', 'dot',
     vdot : Complex-conjugating dot product.
     tensordot : Sum products over arbitrary axes.
     einsum : Einstein summation convention.
+    matmul : '@' operator as method with out parameter.
 
     Examples
     --------
@@ -1954,7 +1955,7 @@ add_newdoc('numpy.core', 'dot',
     >>> np.dot([2j, 3j], [2j, 3j])
     (-13+0j)
 
-    For 2-D arrays it's the matrix product:
+    For 2-D arrays it is the matrix product:
 
     >>> a = [[1, 0], [0, 1]]
     >>> b = [[4, 1], [2, 2]]
@@ -1971,6 +1972,130 @@ add_newdoc('numpy.core', 'dot',
 
     """)
 
+add_newdoc('numpy.core', 'matmul',
+    """
+    matmul(a, b, out=None)
+
+    Matrix product of two arrays.
+
+    The behavior depends on the arguments in the following way.
+
+    - If both arguments are 2-D they are multiplied like conventional
+      matrices.
+    - If either argument is N-D, N > 2, it is treated as a stack of
+      matrices residing in the last two indexes and broadcast accordingly.
+    - If the first argument is 1-D, it is promoted to a matrix by
+      prepending a 1 to its dimensions. After matrix multiplication
+      the prepended 1 is removed.
+    - If the second argument is 1-D, it is promoted to a matrix by
+      appending a 1 to its dimensions. After matrix multiplication
+      the appended 1 is removed.
+
+    Multiplication by a scalar is not allowed, use ``*`` instead. Note that
+    multiplying a stack of matrices with a vector will result in a stack of
+    vectors, but matmul will not recognize it as such.
+
+    ``matmul`` differs from ``dot`` in two important ways.
+
+    - Multiplication by scalars is not allowed.
+    - Stacks of matrices are broadcast together as if the matrices
+      were elements.
+
+    .. warning::
+       This function is preliminary and included in Numpy 1.10 for testing
+       and documentation. Its semantics will not change, but the number and
+       order of the optional arguments will.
+
+    .. versionadded:: 1.10.0
+
+    Parameters
+    ----------
+    a : array_like
+        First argument.
+    b : array_like
+        Second argument.
+    out : ndarray, optional
+        Output argument. This must have the exact kind that would be returned
+        if it was not used. In particular, it must have the right type, must be
+        C-contiguous, and its dtype must be the dtype that would be returned
+        for `dot(a,b)`. This is a performance feature. Therefore, if these
+        conditions are not met, an exception is raised, instead of attempting
+        to be flexible.
+
+    Returns
+    -------
+    output : ndarray
+        Returns the dot product of `a` and `b`.  If `a` and `b` are both
+        1-D arrays then a scalar is returned; otherwise an array is
+        returned.  If `out` is given, then it is returned.
+
+    Raises
+    ------
+    ValueError
+        If the last dimension of `a` is not the same size as
+        the second-to-last dimension of `b`.
+
+        If scalar value is passed.
+
+    See Also
+    --------
+    vdot : Complex-conjugating dot product.
+    tensordot : Sum products over arbitrary axes.
+    einsum : Einstein summation convention.
+    dot : alternative matrix product with different broadcasting rules.
+
+    Notes
+    -----
+    The matmul function implements the semantics of the `@` operator introduced
+    in Python 3.5 following PEP465.
+
+    Examples
+    --------
+    For 2-D arrays it is the matrix product:
+
+    >>> a = [[1, 0], [0, 1]]
+    >>> b = [[4, 1], [2, 2]]
+    >>> np.matmul(a, b)
+    array([[4, 1],
+           [2, 2]])
+
+    For 2-D mixed with 1-D, the result is the usual.
+
+    >>> a = [[1, 0], [0, 1]]
+    >>> b = [1, 2]
+    >>> np.matmul(a, b)
+    array([1, 2])
+    >>> np.matmul(b, a)
+    array([1, 2])
+
+
+    Broadcasting is conventional for stacks of arrays
+
+    >>> a = np.arange(2*2*4).reshape((2,2,4))
+    >>> b = np.arange(2*2*4).reshape((2,4,2))
+    >>> np.matmul(a,b).shape
+    (2, 2, 2)
+    >>> np.matmul(a,b)[0,1,1]
+    98
+    >>> sum(a[0,1,:] * b[0,:,1])
+    98
+
+    Vector, vector returns the scalar inner product, but neither argument
+    is complex-conjugated:
+
+    >>> np.matmul([2j, 3j], [2j, 3j])
+    (-13+0j)
+
+    Scalar multiplication raises an error.
+
+    >>> np.matmul([1,2], 3)
+    Traceback (most recent call last):
+    ...
+    ValueError: Scalar operands are not allowed, use '*' instead
+
+    """)
+
+
 add_newdoc('numpy.core', 'einsum',
     """
     einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='safe')
diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index ea2d4d0a2..bf22f6954 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -43,7 +43,8 @@ __all__ = ['newaxis', 'ndarray', 'flatiter', 'nditer', 'nested_iters', 'ufunc',
            'Inf', 'inf', 'infty', 'Infinity',
            'nan', 'NaN', 'False_', 'True_', 'bitwise_not',
            'CLIP', 'RAISE', 'WRAP', 'MAXDIMS', 'BUFSIZE', 'ALLOW_THREADS',
-           'ComplexWarning', 'may_share_memory', 'full', 'full_like']
+           'ComplexWarning', 'may_share_memory', 'full', 'full_like',
+           'matmul']
 
 if sys.version_info[0] < 3:
     __all__.extend(['getbuffer', 'newbuffer'])
@@ -390,6 +391,11 @@ lexsort = multiarray.lexsort
 compare_chararrays = multiarray.compare_chararrays
 putmask = multiarray.putmask
 einsum = multiarray.einsum
+dot = multiarray.dot
+inner = multiarray.inner
+vdot = multiarray.vdot
+matmul = multiarray.matmul
+
 
 def asarray(a, dtype=None, order=None):
     """
@@ -1081,11 +1087,6 @@ def outer(a, b, out=None):
     b = asarray(b)
     return multiply(a.ravel()[:, newaxis], b.ravel()[newaxis,:], out)
 
-# try to import blas optimized dot if available
-envbak = os.environ.copy()
-dot = multiarray.dot
-inner = multiarray.inner
-vdot = multiarray.vdot
 
 def alterdot():
     """
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 369b5a8e1..7980a0de7 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -27,8 +27,8 @@
 #include "numpy/npy_math.h"
 
 #include "npy_config.h"
-
 #include "npy_pycompat.h"
+#include "npy_import.h"
 
 NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0;
 
@@ -64,6 +64,7 @@ NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0;
 /* Only here for API compatibility */
 NPY_NO_EXPORT PyTypeObject PyBigArray_Type;
 
+
 /*NUMPY_API
  * Get Priority from object
  */
@@ -239,7 +240,7 @@ PyArray_AsCArray(PyObject **op, void *ptr, npy_intp *dims, int nd,
     *op = (PyObject *)ap;
     return 0;
 
- fail:
+fail:
     PyErr_SetString(PyExc_MemoryError, "no memory");
     return -1;
 }
@@ -930,7 +931,7 @@ PyArray_InnerProduct(PyObject *op1, PyObject *op2)
     Py_DECREF(ap2);
     return (PyObject *)ret;
 
- fail:
+fail:
     Py_XDECREF(ap1);
     Py_XDECREF(ap2);
     Py_XDECREF(ret);
@@ -1049,7 +1050,8 @@ PyArray_MatrixProduct2(PyObject *op1, PyObject *op2, PyArrayObject* out)
         goto fail;
     }
 
-    op = PyArray_DATA(ret); os = PyArray_DESCR(ret)->elsize;
+    op = PyArray_DATA(ret);
+    os = PyArray_DESCR(ret)->elsize;
     axis = PyArray_NDIM(ap1)-1;
     it1 = (PyArrayIterObject *)
         PyArray_IterAllButAxis((PyObject *)ap1, &axis);
@@ -1083,7 +1085,7 @@ PyArray_MatrixProduct2(PyObject *op1, PyObject *op2, PyArrayObject* out)
     Py_DECREF(ap2);
     return (PyObject *)ret;
 
- fail:
+fail:
     Py_XDECREF(ap1);
     Py_XDECREF(ap2);
     Py_XDECREF(ret);
@@ -1844,7 +1846,7 @@ array_copyto(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kwds)
 
     Py_RETURN_NONE;
 
- fail:
+fail:
     Py_XDECREF(src);
     Py_XDECREF(wheremask);
     return NULL;
@@ -1887,7 +1889,7 @@ array_empty(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kwds)
     PyDimMem_FREE(shape.ptr);
     return (PyObject *)ret;
 
- fail:
+fail:
     Py_XDECREF(typecode);
     PyDimMem_FREE(shape.ptr);
     return NULL;
@@ -1918,7 +1920,7 @@ array_empty_like(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kwds)
 
     return (PyObject *)ret;
 
- fail:
+fail:
     Py_XDECREF(prototype);
     Py_XDECREF(dtype);
     return NULL;
@@ -2041,7 +2043,7 @@ array_zeros(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kwds)
     PyDimMem_FREE(shape.ptr);
     return (PyObject *)ret;
 
- fail:
+fail:
     Py_XDECREF(typecode);
     PyDimMem_FREE(shape.ptr);
     return (PyObject *)ret;
@@ -2369,6 +2371,170 @@ fail:
 }
 
 
+
+/*
+ * matmul
+ *
+ * Implements the protocol used by the '@' operator defined in PEP 364.
+ * Not in the NUMPY API at this time, maybe later.
+ *
+ *
+ * in1:        Left hand side operand
+ * in2:        Right hand side operand
+ * out:        Either NULL, or an array into which the output should be placed.
+ *
+ * Returns NULL on error.
+ * Returns NotImplemented on priority override.
+ */
+static PyObject *
+array_matmul(PyObject *NPY_UNUSED(m), PyObject *args, PyObject* kwds)
+{
+    static PyObject *matmul = NULL;
+    int errval;
+    PyObject *override = NULL;
+    PyObject *in1, *in2, *out = NULL;
+    char* kwlist[] = {"a", "b", "out", NULL };
+    PyArrayObject *ap1, *ap2, *ret = NULL;
+    NPY_ORDER order = NPY_KEEPORDER;
+    NPY_CASTING casting = NPY_SAFE_CASTING;
+    PyArray_Descr *dtype;
+    int nd1, nd2, typenum;
+    char *subscripts;
+    PyArrayObject *ops[2];
+
+    npy_cache_pyfunc("numpy.core.multiarray", "matmul", &matmul);
+    if (matmul == NULL) {
+        return NULL;
+    }
+
+    errval = PyUFunc_CheckOverride(matmul, "__call__",
+                                   args, kwds, &override, 2);
+    if (errval) {
+        return NULL;
+    }
+    else if (override) {
+        return override;
+    }
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|O", kwlist,
+                                     &in1, &in2, &out)) {
+        return NULL;
+    }
+
+    if (out == Py_None) {
+        out = NULL;
+    }
+    if (out != NULL && !PyArray_Check(out)) {
+        PyErr_SetString(PyExc_TypeError,
+                        "'out' must be an array");
+        return NULL;
+    }
+
+    dtype = PyArray_DescrFromObject(in1, NULL);
+    dtype = PyArray_DescrFromObject(in2, dtype);
+    if (dtype == NULL) {
+        PyErr_SetString(PyExc_ValueError,
+                "Cannot find a common data type.");
+        return NULL;
+    }
+    typenum = dtype->type_num;
+
+    if (typenum == NPY_OBJECT) {
+        /* matmul is not currently implemented for object arrays */
+        PyErr_SetString(PyExc_TypeError,
+                "Object arrays are not currently supported");
+        Py_DECREF(dtype);
+        return NULL;
+    }
+
+    ap1 = (PyArrayObject *)PyArray_FromAny(in1, dtype, 0, 0,
+                                           NPY_ARRAY_ALIGNED, NULL);
+    if (ap1 == NULL) {
+        return NULL;
+    }
+
+    Py_INCREF(dtype);
+    ap2 = (PyArrayObject *)PyArray_FromAny(in2, dtype, 0, 0,
+                                           NPY_ARRAY_ALIGNED, NULL);
+    if (ap2 == NULL) {
+        Py_DECREF(ap1);
+        return NULL;
+    }
+
+    if (PyArray_NDIM(ap1) == 0 || PyArray_NDIM(ap2) == 0) {
+        /* Scalars are rejected */
+        PyErr_SetString(PyExc_ValueError,
+                "Scalar operands are not allowed, use '*' instead");
+        return NULL;
+    }
+
+    nd1 = PyArray_NDIM(ap1);
+    nd2 = PyArray_NDIM(ap2);
+
+#if defined(HAVE_CBLAS)
+    if (nd1 <= 2 && nd2 <= 2 &&
+            (NPY_DOUBLE == typenum || NPY_CDOUBLE == typenum ||
+             NPY_FLOAT == typenum || NPY_CFLOAT == typenum)) {
+        return cblas_matrixproduct(typenum, ap1, ap2, out);
+    }
+#endif
+
+    /*
+     * Use einsum for the stacked cases. This is a quick implementation
+     * to avoid setting up the proper iterators. Einsum broadcasts, so
+     * we need to check dimensions before the call.
+     */
+    if (nd1 == 1 && nd2 == 1) {
+        /* vector vector */
+        if (PyArray_DIM(ap1, 0) != PyArray_DIM(ap2, 0)) {
+            dot_alignment_error(ap1, 0, ap2, 0);
+            goto fail;
+        }
+        subscripts = "i, i";
+    }
+    else if (nd1 == 1) {
+        /* vector  matrix */
+        if (PyArray_DIM(ap1, 0) != PyArray_DIM(ap2, nd2 - 2)) {
+            dot_alignment_error(ap1, 0, ap2, nd2 - 2);
+            goto fail;
+        }
+        subscripts = "i, ...ij";
+    }
+    else if (nd2 == 1) {
+        /* matrix  vector */
+        if (PyArray_DIM(ap1, nd1 - 1) != PyArray_DIM(ap2, 0)) {
+            dot_alignment_error(ap1, nd1 - 1, ap2, 0);
+            goto fail;
+        }
+        subscripts = "...i, i";
+    }
+    else {
+        /* matrix * matrix */
+        if (PyArray_DIM(ap1, nd1 - 1) != PyArray_DIM(ap2, nd2 - 2)) {
+            dot_alignment_error(ap1, nd1 - 1, ap2, nd2 - 2);
+            goto fail;
+        }
+        subscripts = "...ij, ...jk";
+    }
+    ops[0] = ap1;
+    ops[1] = ap2;
+    ret = PyArray_EinsteinSum(subscripts, 2, ops, NULL, order, casting, out);
+    Py_DECREF(ap1);
+    Py_DECREF(ap2);
+
+    /* If no output was supplied, possibly convert to a scalar */
+    if (ret != NULL && out == NULL) {
+        ret = PyArray_Return((PyArrayObject *)ret);
+    }
+    return (PyObject *)ret;
+
+fail:
+    Py_XDECREF(ap1);
+    Py_XDECREF(ap2);
+    return NULL;
+}
+
+
 static int
 einsum_sub_op_from_str(PyObject *args, PyObject **str_obj, char **subscripts,
                        PyArrayObject **op)
@@ -2862,7 +3028,7 @@ array__reconstruct(PyObject *NPY_UNUSED(dummy), PyObject *args)
 
     return ret;
 
- fail:
+fail:
     evil_global_disable_warn_O4O8_flag = 0;
 
     Py_XDECREF(dtype);
@@ -3090,7 +3256,7 @@ PyArray_Where(PyObject *condition, PyObject *x, PyObject *y)
         return ret;
     }
 
- fail:
+fail:
     Py_DECREF(arr);
     Py_XDECREF(ax);
     Py_XDECREF(ay);
@@ -3936,6 +4102,9 @@ static struct PyMethodDef array_module_methods[] = {
     {"vdot",
         (PyCFunction)array_vdot,
         METH_VARARGS | METH_KEYWORDS, NULL},
+    {"matmul",
+        (PyCFunction)array_matmul,
+        METH_VARARGS | METH_KEYWORDS, NULL},
     {"einsum",
         (PyCFunction)array_einsum,
         METH_VARARGS|METH_KEYWORDS, NULL},
diff --git a/numpy/core/src/multiarray/number.c b/numpy/core/src/multiarray/number.c
index 168799f11..a2160afd8 100644
--- a/numpy/core/src/multiarray/number.c
+++ b/numpy/core/src/multiarray/number.c
@@ -8,11 +8,10 @@
 #include "numpy/arrayobject.h"
 
 #include "npy_config.h"
-
 #include "npy_pycompat.h"
-
-#include "number.h"
+#include "npy_import.h"
 #include "common.h"
+#include "number.h"
 
 /*************************************************************************
  ****************   Implement Number Protocol ****************************
@@ -386,6 +385,24 @@ array_remainder(PyArrayObject *m1, PyObject *m2)
     return PyArray_GenericBinaryFunction(m1, m2, n_ops.remainder);
 }
 
+
+#if PY_VERSION_HEX >= 0x03050000
+/* Need this to be version dependent on account of the slot check */
+static PyObject *
+array_matrix_multiply(PyArrayObject *m1, PyObject *m2)
+{
+    static PyObject *matmul = NULL;
+
+    npy_cache_pyfunc("numpy.core.multiarray", "matmul", &matmul);
+    if (matmul == NULL) {
+        return NULL;
+    }
+    GIVE_UP_IF_HAS_RIGHT_BINOP(m1, m2, "__matmul__", "__rmatmul__",
+                               0, nb_matrix_multiply);
+    return PyArray_GenericBinaryFunction(m1, m2, matmul);
+}
+#endif
+
 /* Determine if object is a scalar and if so, convert the object
  *   to a double and place it in the out_exponent argument
  *   and return the "scalar kind" as a result.   If the object is
@@ -723,6 +740,7 @@ array_inplace_true_divide(PyArrayObject *m1, PyObject *m2)
                                                 n_ops.true_divide);
 }
 
+
 static int
 _array_nonzero(PyArrayObject *mp)
 {
@@ -1066,5 +1084,9 @@ NPY_NO_EXPORT PyNumberMethods array_as_number = {
     (binaryfunc)array_true_divide,              /*nb_true_divide*/
     (binaryfunc)array_inplace_floor_divide,     /*nb_inplace_floor_divide*/
     (binaryfunc)array_inplace_true_divide,      /*nb_inplace_true_divide*/
-    (unaryfunc)array_index,                     /* nb_index */
+    (unaryfunc)array_index,                     /*nb_index */
+#if PY_VERSION_HEX >= 0x03050000
+    (binaryfunc)array_matrix_multiply,          /*nb_matrix_multiply*/
+    (binaryfunc)NULL,                           /*nb_inplacematrix_multiply*/
+#endif
 };
diff --git a/numpy/core/src/private/npy_import.h b/numpy/core/src/private/npy_import.h
index e618e121e..a75c59884 100644
--- a/numpy/core/src/private/npy_import.h
+++ b/numpy/core/src/private/npy_import.h
@@ -2,7 +2,6 @@
 #define NPY_IMPORT_H
 
 #include <Python.h>
-#include <assert.h>
 
 /*! \brief Fetch and cache Python function.
  *
@@ -17,7 +16,7 @@
  * @param function Function name.
  * @param cache Storage location for imported function.
  */
-NPY_INLINE void
+NPY_INLINE static void
 npy_cache_pyfunc(const char *module, const char *function, PyObject **cache)
 {
     if (*cache == NULL) {
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 33b75d490..f37668de9 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -8,6 +8,7 @@ import shutil
 import warnings
 import operator
 import io
+import itertools
 if sys.version_info[0] >= 3:
     import builtins
 else:
@@ -3927,7 +3928,6 @@ class TestDot(TestCase):
         assert_raises(TypeError, c.dot, b)
 
     def test_accelerate_framework_sgemv_fix(self):
-        from itertools import product
         if sys.platform != 'darwin':
             return
 
@@ -3954,7 +3954,7 @@ class TestDot(TestCase):
         s = aligned_array((100, 100), 15, np.float32)
         np.dot(s, m) # this will always segfault if the bug is present
 
-        testdata = product((15,32), (10000,), (200,89), ('C','F'))
+        testdata = itertools.product((15,32), (10000,), (200,89), ('C','F'))
         for align, m, n, a_order in testdata:
             # Calculation in double precision
             A_d = np.random.rand(m, n)
@@ -3994,6 +3994,287 @@ class TestDot(TestCase):
             assert_dot_close(A_f_12, X_f_2, desired)
 
 
+class MatmulCommon():
+    """Common tests for '@' operator and numpy.matmul.
+
+    Do not derive from TestCase to avoid nose running it.
+
+    """
+    # Should work with these types. Will want to add
+    # "O" at some point
+    types = "?bhilqBHILQefdgFDG"
+
+    def test_exceptions(self):
+        dims = [
+            ((1,), (2,)),            # mismatched vector vector
+            ((2, 1,), (2,)),         # mismatched matrix vector
+            ((2,), (1, 2)),          # mismatched vector matrix
+            ((1, 2), (3, 1)),        # mismatched matrix matrix
+            ((1,), ()),              # vector scalar
+            ((), (1)),               # scalar vector
+            ((1, 1), ()),            # matrix scalar
+            ((), (1, 1)),            # scalar matrix
+            ((2, 2, 1), (3, 1, 2)),  # cannot broadcast
+            ]
+
+        for dt, (dm1, dm2) in itertools.product(self.types, dims):
+            a = np.ones(dm1, dtype=dt)
+            b = np.ones(dm2, dtype=dt)
+            assert_raises(ValueError, self.matmul, a, b)
+
+    def test_shapes(self):
+        dims = [
+            ((1, 1), (2, 1, 1)),     # broadcast first argument
+            ((2, 1, 1), (1, 1)),     # broadcast second argument
+            ((2, 1, 1), (2, 1, 1)),  # matrix stack sizes match
+            ]
+
+        for dt, (dm1, dm2) in itertools.product(self.types, dims):
+            a = np.ones(dm1, dtype=dt)
+            b = np.ones(dm2, dtype=dt)
+            res = self.matmul(a, b)
+            assert_(res.shape == (2, 1, 1))
+
+        # vector vector returns scalars.
+        for dt in self.types:
+            a = np.ones((2,), dtype=dt)
+            b = np.ones((2,), dtype=dt)
+            c = self.matmul(a, b)
+            assert_(np.array(c).shape == ())
+
+    def test_result_types(self):
+        mat = np.ones((1,1))
+        vec = np.ones((1,))
+        for dt in self.types:
+            m = mat.astype(dt)
+            v = vec.astype(dt)
+            for arg in [(m, v), (v, m), (m, m)]:
+                res = matmul(*arg)
+                assert_(res.dtype == dt)
+
+            # vector vector returns scalars
+            res = matmul(v, v)
+            assert_(type(res) is dtype(dt).type)
+
+    def test_vector_vector_values(self):
+        vec = np.array([1, 2])
+        tgt = 5
+        for dt in self.types[1:]:
+            v1 = vec.astype(dt)
+            res = matmul(v1, v1)
+            assert_equal(res, tgt)
+
+        # boolean type
+        vec = np.array([True, True], dtype='?')
+        res = matmul(vec, vec)
+        assert_equal(res, True)
+
+    def test_vector_matrix_values(self):
+        vec = np.array([1, 2])
+        mat1 = np.array([[1, 2], [3, 4]])
+        mat2 = np.stack([mat1]*2, axis=0)
+        tgt1 = np.array([7, 10])
+        tgt2 = np.stack([tgt1]*2, axis=0)
+        for dt in self.types[1:]:
+            v = vec.astype(dt)
+            m1 = mat1.astype(dt)
+            m2 = mat2.astype(dt)
+            res = matmul(v, m1)
+            assert_equal(res, tgt1)
+            res = matmul(v, m2)
+            assert_equal(res, tgt2)
+
+        # boolean type
+        vec = np.array([True, False])
+        mat1 = np.array([[True, False], [False, True]])
+        mat2 = np.stack([mat1]*2, axis=0)
+        tgt1 = np.array([True, False])
+        tgt2 = np.stack([tgt1]*2, axis=0)
+
+        res = matmul(vec, mat1)
+        assert_equal(res, tgt1)
+        res = matmul(vec, mat2)
+        assert_equal(res, tgt2)
+
+    def test_matrix_vector_values(self):
+        vec = np.array([1, 2])
+        mat1 = np.array([[1, 2], [3, 4]])
+        mat2 = np.stack([mat1]*2, axis=0)
+        tgt1 = np.array([5, 11])
+        tgt2 = np.stack([tgt1]*2, axis=0)
+        for dt in self.types[1:]:
+            v = vec.astype(dt)
+            m1 = mat1.astype(dt)
+            m2 = mat2.astype(dt)
+            res = matmul(m1, v)
+            assert_equal(res, tgt1)
+            res = matmul(m2, v)
+            assert_equal(res, tgt2)
+
+        # boolean type
+        vec = np.array([True, False])
+        mat1 = np.array([[True, False], [False, True]])
+        mat2 = np.stack([mat1]*2, axis=0)
+        tgt1 = np.array([True, False])
+        tgt2 = np.stack([tgt1]*2, axis=0)
+
+        res = matmul(vec, mat1)
+        assert_equal(res, tgt1)
+        res = matmul(vec, mat2)
+        assert_equal(res, tgt2)
+
+    def test_matrix_matrix_values(self):
+        mat1 = np.array([[1, 2], [3, 4]])
+        mat2 = np.array([[1, 0], [1, 1]])
+        mat12 = np.stack([mat1, mat2], axis=0)
+        mat21 = np.stack([mat2, mat1], axis=0)
+        tgt11 = np.array([[7, 10], [15, 22]])
+        tgt12 = np.array([[3, 2], [7, 4]])
+        tgt21 = np.array([[1, 2], [4, 6]])
+        tgt22 = np.array([[1, 0], [2, 1]])
+        tgt12_21 = np.stack([tgt12, tgt21], axis=0)
+        tgt11_12 = np.stack((tgt11, tgt12), axis=0)
+        tgt11_21 = np.stack((tgt11, tgt21), axis=0)
+        for dt in self.types[1:]:
+            m1 = mat1.astype(dt)
+            m2 = mat2.astype(dt)
+            m12 = mat12.astype(dt)
+            m21 = mat21.astype(dt)
+
+            # matrix @ matrix
+            res = matmul(m1, m2)
+            assert_equal(res, tgt12)
+            res = matmul(m2, m1)
+            assert_equal(res, tgt21)
+
+            # stacked @ matrix
+            res = self.matmul(m12, m1)
+            assert_equal(res, tgt11_21)
+
+            # matrix @ stacked
+            res = self.matmul(m1, m12)
+            assert_equal(res, tgt11_12)
+
+            # stacked @ stacked
+            res = self.matmul(m12, m21)
+            assert_equal(res, tgt12_21)
+
+        # boolean type
+        m1 = np.array([[1, 1], [0, 0]], dtype=np.bool_)
+        m2 = np.array([[1, 0], [1, 1]], dtype=np.bool_)
+        m12 = np.stack([m1, m2], axis=0)
+        m21 = np.stack([m2, m1], axis=0)
+        tgt11 = m1
+        tgt12 = m1
+        tgt21 = np.array([[1, 1], [1, 1]], dtype=np.bool_)
+        tgt22 = m2
+        tgt12_21 = np.stack([tgt12, tgt21], axis=0)
+        tgt11_12 = np.stack((tgt11, tgt12), axis=0)
+        tgt11_21 = np.stack((tgt11, tgt21), axis=0)
+
+        # matrix @ matrix
+        res = matmul(m1, m2)
+        assert_equal(res, tgt12)
+        res = matmul(m2, m1)
+        assert_equal(res, tgt21)
+
+        # stacked @ matrix
+        res = self.matmul(m12, m1)
+        assert_equal(res, tgt11_21)
+
+        # matrix @ stacked
+        res = self.matmul(m1, m12)
+        assert_equal(res, tgt11_12)
+
+        # stacked @ stacked
+        res = self.matmul(m12, m21)
+        assert_equal(res, tgt12_21)
+
+    def test_numpy_ufunc_override(self):
+
+        class A(np.ndarray):
+            def __new__(cls, *args, **kwargs):
+                return np.array(*args, **kwargs).view(cls)
+            def __numpy_ufunc__(self, ufunc, method, pos, inputs, **kwargs):
+                return "A"
+
+        class B(np.ndarray):
+            def __new__(cls, *args, **kwargs):
+                return np.array(*args, **kwargs).view(cls)
+            def __numpy_ufunc__(self, ufunc, method, pos, inputs, **kwargs):
+                return NotImplemented
+
+        a = A([1, 2])
+        b = B([1, 2])
+        c = ones(2)
+        assert_equal(self.matmul(a, b), "A")
+        assert_equal(self.matmul(b, a), "A")
+        assert_raises(TypeError, self.matmul, b, c)
+
+
+class TestMatmul(MatmulCommon, TestCase):
+    matmul = np.matmul
+
+    def test_out_arg(self):
+        a = np.ones((2, 2), dtype=np.float)
+        b = np.ones((2, 2), dtype=np.float)
+        tgt = np.full((2,2), 2, dtype=np.float)
+
+        # test as positional argument
+        msg = "out positional argument"
+        out = np.zeros((2, 2), dtype=np.float)
+        self.matmul(a, b, out)
+        assert_array_equal(out, tgt, err_msg=msg)
+
+        # test as keyword argument
+        msg = "out keyword argument"
+        out = np.zeros((2, 2), dtype=np.float)
+        self.matmul(a, b, out=out)
+        assert_array_equal(out, tgt, err_msg=msg)
+
+        # test out with not allowed type cast (safe casting)
+        # einsum and cblas raise different error types, so
+        # use Exception.
+        msg = "out argument with illegal cast"
+        out = np.zeros((2, 2), dtype=np.int32)
+        assert_raises(Exception, self.matmul, a, b, out=out)
+
+        # skip following tests for now, cblas does not allow non-contiguous
+        # outputs and consistency with dot would require same type,
+        # dimensions, subtype, and c_contiguous.
+
+        # test out with allowed type cast
+        # msg = "out argument with allowed cast"
+        # out = np.zeros((2, 2), dtype=np.complex128)
+        # self.matmul(a, b, out=out)
+        # assert_array_equal(out, tgt, err_msg=msg)
+
+        # test out non-contiguous
+        # msg = "out argument with non-contiguous layout"
+        # c = np.zeros((2, 2, 2), dtype=np.float)
+        # self.matmul(a, b, out=c[..., 0])
+        # assert_array_equal(c, tgt, err_msg=msg)
+
+
+if sys.version_info[:2] >= (3, 5):
+    class TestMatmulOperator(MatmulCommon, TestCase):
+        from operator import matmul
+
+        def test_array_priority_override(self):
+
+            class A(object):
+                __array_priority__ = 1000
+                def __matmul__(self, other):
+                    return "A"
+                def __rmatmul__(self, other):
+                    return "A"
+
+            a = A()
+            b = ones(2)
+            assert_equal(self.matmul(a, b), "A")
+            assert_equal(self.matmul(b, a), "A")
+
+
 class TestInner(TestCase):
 
     def test_inner_scalar_and_matrix_of_objects(self):