Add support for Python's matrix multiplication operator

Author: bjodah

symengine/lib/symengine_wrapper.pyx

@@ -3271,6 +3271,13 @@ cdef class DenseMatrixBase(MatrixBase):
         else:
             return NotImplemented
 
+    def __matmul__(a, b):
+        a = _sympify(a, False)
+        b = _sympify(b, False)
+        if (a.ncols() != b.nrows()):
+            raise ShapeError("Invalid shapes for matrix multiplication. Got %s %s" % (a.shape, b.shape))
+        return a.mul_matrix(b)
+
     def __truediv__(a, b):
         return div_matrices(a, b)
 
@@ -4889,7 +4896,7 @@ cdef class LambdaDouble(_Lambdify):
         c_inp = np.ascontiguousarray(inp.ravel(order=self.order), dtype=self.numpy_dtype)
         c_out = out
         for idx in range(nbroadcast):
-            self.lambda_double[0].call(&c_out[idx*self.tot_out_size], &c_inp[idx*self.args_size]) 
+            self.lambda_double[0].call(&c_out[idx*self.tot_out_size], &c_inp[idx*self.args_size])
 
     cpdef as_scipy_low_level_callable(self):
         from ctypes import c_double, c_void_p, c_int, cast, POINTER, CFUNCTYPE
@@ -5135,7 +5142,7 @@ def Lambdify(args, *exprs, cppbool real=True, backend=None, order='C',
                     raise ValueError("Long double not supported on this platform")
             else:
                 raise ValueError("Unknown numpy dtype.")
-                
+
             if as_scipy:
                 return ret.as_scipy_low_level_callable()
             return ret

symengine/tests/test_matrices.py

@@ -229,6 +229,8 @@ def test_mul_matrix():
 
     assert A.mul_matrix(B) == DenseMatrix(2, 2, [a + b, 0, c + d, 0])
     assert A * B == DenseMatrix(2, 2, [a + b, 0, c + d, 0])
+    assert A @ B == DenseMatrix(2, 2, [a + b, 0, c + d, 0])
+    assert (A @ DenseMatrix(2, 1, [0]*4)).shape == (2, 1)
 
     C = DenseMatrix(2, 3, [1, 2, 3, 2, 3, 4])
     D = DenseMatrix(3, 2, [3, 4, 4, 5, 5, 6])