Simplify _infer_size

src/Coloring/Coloring.jl

@@ -30,14 +30,14 @@ IndexedSet(n::Integer) = IndexedSet(zeros(Int, n), trues(n), 0)
 
 function Base.push!(v::IndexedSet, i::Integer)
     if v.empty[i]  # new index
-        v.nzidx[v.nnz+=1] = i
+        v.nzidx[v.nnz += 1] = i
         v.empty[i] = false
     end
     return
 end
 
 function Base.empty!(v::IndexedSet)
-    for i in 1:v.nnz
+    for i in 1:(v.nnz)
         v.empty[v.nzidx[i]] = true
     end
     v.nnz = 0
@@ -58,7 +58,7 @@ function Base.resize!(v::IndexedSet, n::Integer)
     return
 end
 
-Base.collect(v::IndexedSet) = v.nzidx[1:v.nnz]
+Base.collect(v::IndexedSet) = v.nzidx[1:(v.nnz)]
 
 function Base.union!(v::IndexedSet, s)
     for x in s

src/reverse_mode.jl

@@ -133,7 +133,7 @@ function _forward_eval(
             len = length(tape_range)
             copyto!(
                 view(f.forward_storage, tape_range),
-                view(x, node.index:(node.index+len-1)),
+                view(x, (node.index):(node.index+len-1)),
             )
         elseif node.type == NODE_VALUE_BLOCK
             # Pre-loaded into `forward_storage` at construction.
@@ -745,6 +745,16 @@ function __reverse_broadcasted_div(f, ilhs, irhs, dout, dlhs, drhs)
     )
 end
 
+# Reverse for `:sum_dims`. `y = sum(x; dims=d)` collapses one axis, so
+# ∂y[i…]/∂x[j…] is 1 iff `j` matches `i` on every non-reduced axis (any
+# `j` along the reduced axis maps to the same `y` slot). Broadcasting the
+# (m,1) or (1,n) parent adjoint across the full child shape produces
+# exactly that pattern.
+function _reverse_sum_dims!(rev_arr, rev_parent)
+    rev_arr .= rev_parent
+    return
+end
+
 """
     _reverse_eval(f::_SubexpressionStorage)
 
@@ -1169,7 +1179,7 @@ function _extract_reverse_pass_inner(
         if node.type == NODE_VARIABLE_BLOCK
             tape_range = _storage_range(f.sizes, k)
             len = length(tape_range)
-            x_range = node.index:(node.index+len-1)
+            x_range = (node.index):(node.index+len-1)
             view(output, x_range) .+=
                 scale .* view(f.reverse_storage, tape_range)
         elseif node.type == NODE_VARIABLE