ndarray: support matrix/tensor multiplication (#384)

NEWS.md

@@ -161,6 +161,26 @@
    1.0  2.0  3.0  4.0
   ```
 
+* Matrix/tensor multiplication is supported now. (#TBD)
+
+  ```julia
+  julia> x
+  2×3 mx.NDArray{Float32,2} @ CPU0:
+   1.0  2.0  3.0
+   4.0  5.0  6.0
+
+  julia> y
+  3 mx.NDArray{Float32,1} @ CPU0:
+   -1.0
+   -2.0
+   -3.0
+
+  julia> x * y
+  2 mx.NDArray{Float32,1} @ CPU0:
+   -14.0
+   -32.0
+  ```
+
 ## API Changes
 
 ### `NDArray`

src/ndarray.jl

@@ -674,7 +674,7 @@ import Base: *
 """
     .*(x, y)
 
-Currently only multiplication a scalar with an `NDArray` is implemented.
+Elementwise multiplication for `NDArray`.
 """
 *(x::NDArray, y::Real)  = _mul_scalar(x, scalar = y)
 *(y::Real, x::NDArray)  = _mul_scalar(x, scalar = y)
@@ -686,9 +686,9 @@ broadcast_(::typeof(*), x::NDArray, y::NDArray) = _mul(x, y)
 """
     *(A::NDArray, B::NDArray)
 
-Matrix (2D NDArray) multiplication.
+Matrix/tensor multiplication.
 """
-*(x::NDArray{T,2}, y::NDArray{S,2}) where {T,S} = dot(x, y)
+*(x::NDArray{T}, y::NDArray{T}) where T = x ⋅ y
 
 """
     div_from!(dst::NDArray, arg::NDArrayOrReal)
@@ -1165,7 +1165,7 @@ _mxsig[:reshape] = :(reshape(arr; shape = dim, reverse = !reverse))
 @_remap minimum(arr::NDArray, dims) min(arr; axis = 0 .- dims, keepdims = true)
 
 # See https://github.com/dmlc/MXNet.jl/issues/55
-@_remap dot(x::NDArray{T,N}, y::NDArray{S,N}) where {T,S,N} dot(y, x)
+@_remap dot(x::NDArray, y::NDArray) dot(y, x)
 
 # See https://github.com/dmlc/MXNet.jl/pull/123
 @_remap transpose(arr::NDArray{T,1}) where T reshape(arr; shape = (1, length(arr)), reverse = true)

test/unittest/ndarray.jl

@@ -740,7 +740,28 @@ function test_dot()
 
   x = mx.zeros(1, 2)
   y = mx.zeros(1, 2, 3)
-  @test_throws MethodError dot(x, y)
+  @test_throws mx.MXError dot(x, y)  # dimension mismatch
+
+  info("NDArray::matrix mul")
+  let
+    A = [1. 2 3; 4 5 6]
+    B = [-1., -2, -3]
+    x = NDArray(A)
+    y = NDArray(B)
+    z = x * y
+    @test copy(z) == A * B
+    @test size(z) == (2,)
+  end
+
+  let
+    A = [1. 2 3; 4 5 6]
+    B = [-1. -2; -3 -4; -5 -6]
+    x = NDArray(A)
+    y = NDArray(B)
+    z = x * y
+    @test copy(z) == A * B
+    @test size(z) == (2, 2)
+  end
 end
 
 function test_eltype()