From d8937fbd5273aa66c89bb932b428e236d1376d15 Mon Sep 17 00:00:00 2001
From: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
Date: Tue, 21 Sep 2021 16:36:26 -0400
Subject: [PATCH] Update src/lib/broadcast.jl

---
 src/lib/broadcast.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/lib/broadcast.jl b/src/lib/broadcast.jl
index 33274710a..1f3ec0d8f 100644
--- 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