Merge pull request #19 from ToucheSir/bc/fix-rank

Report correct rank for matmul and aggregrations
FluxML · May 12, 2020 · 72c0dde · 72c0dde
2 parents eb8f25c + 7149c0e
commit 72c0dde
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/src/ops.jl b/src/ops.jl
@@ -3,9 +3,11 @@ import Base: +, -, *, /
 for (op,fn) in zip((:+, :-, :/, :*), (atg_add, atg_sub, atg_div, atg_matmul))
   @eval function $op(t1::Tensor{T,N}, t2::Tensor{T,K}) where {T,N,K}
     ptr = Ref(Ptr{Cvoid}())
+    rank = Ref{Cint}(-1)
 
     $fn(ptr, t1.ptr, t2.ptr)
-    Tensor{T,N}(ptr[], on(t1))
+    at_dim(rank, ptr[])
+    Tensor{T,rank[]}(ptr[], on(t1))
   end
 end
 

diff --git a/src/statistics.jl b/src/statistics.jl
@@ -5,21 +5,21 @@ function Statistics.mean(t::Tensor{T,N}; dims = :) where {T,N}
 
   if dims isa Colon
     atg_mean(ptr, t.ptr, options[T])
+    Tensor{T,0}(ptr[], on(t))
   else
     atg_mean1(ptr, t.ptr, dims, length(dims), dims[1], options[T])
+    Tensor{T,N-length(dims)}(ptr[], on(t))
   end
-
-  Tensor{T,N}(ptr[], on(t))
 end
 
 function Statistics.sum(t::Tensor{T,N}; dims = :) where {T,N}
   ptr = Ref(Ptr{Cvoid}())
 
   if dims isa Colon
     atg_sum(ptr, t.ptr, options[T])
+    Tensor{T,0}(ptr[], on(t))
   else
     atg_sum1(ptr, t.ptr, dims, length(dims), dims[1], options[T])
+    Tensor{T,N-length(dims)}(ptr[], on(t))
   end
-
-  Tensor{T,N}(ptr[], on(t))
 end