Update src/lib/broadcast.jl

FluxML · Sep 22, 2021 · 3bc2e09 · 3bc2e09
1 parent 80123a1
commit 3bc2e09
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/lib/broadcast.jl b/src/lib/broadcast.jl
@@ -50,8 +50,7 @@ function unbroadcast(x::AbstractArray, x̄)
   if length(x) == length(x̄)
     _project(x, x̄)  # ProjectTo handles reshape, offsets, structured matrices, row vectors
   else
-    tup = filter(d -> size(x, d) == 1, ntuple(identity, N))
-    dims = length(tup) == 1 ? first(tup) : tup  # avoid sum(xbar, dims=(1,)) as e.g. sum(SA[1 2; 3 4], dims=(1,)) fails
+    dims = ntuple(d -> size(x, d) == 1 ? d : ndims(x̄)+1, ndims(x̄))
     _project(x, accum_sum(x̄; dims = dims))
   end
 end