fix tests for PR, add comments

Author: Michael Abbott

src/SliceMap.jl

@@ -25,6 +25,8 @@ It provides a gradient for Tracker and Zygote, saving the backward function for
 Any arguments after the matrix are passed to `f` as scalars, i.e.
 `mapcols(f, m, args...) = reduce(hcat, f(col, args...) for col in eeachcol(m))`.
 They do not get sliced/iterated (unlike `map`), nor are their gradients tracked.
+
+Note that if `f` itself contains parameters, their gradients are also not tracked.
 """
 mapcols(f, M, args...) = _mapcols(map, f, M, args...)
 tmapcols(f, M, args...) = _mapcols(threadmap, f, M, args...)
@@ -90,6 +92,9 @@ Like `mapcols()`, but for any slice. The function `f` must preserve shape,
 e.g. if `dims=(2,4)` then `f` must map matrices to matrices.
 
 The gradient is for Zygote only.
+
+Parameters within the function `f` (if there are any) should be correctly tracked,
+which is not the case for `mapcols()`.
 """
 function slicemap(f, A::AbstractArray{T,N}, args...; dims) where {T,N}
     code = ntuple(d -> d in dims ? True() : False(), N)

test/runtests.jl

@@ -146,15 +146,24 @@ end
     xp = Tracker.param(x)
     Tracker.back!(sum(mapcols(F(wp), xp)))
     @test Tracker.grad(xp) ≈ gradx
-    @test_broken Tracker.grad(wp) ≈ gradw # zero
+    @test Tracker.grad(wp) == 0 .* gradw # bug or a feature?
 
-    grad_mapcols = Zygote.gradient(() -> sum(mapcols(F(w), x)), Zygote.Params([w,x]))
+    # fp = F(wp)
+    # wp.grad .= 0; xp.grad .= 0;
+    # Tracker.back!(sum(mapcols(fp, xp)))
+    # @test Tracker.grad(xp) ≈ gradx
+    # @test_broken Tracker.grad(wp) ≈ gradw # zero
+
+    f = F(w)
+    grad_mapcols = Zygote.gradient(() -> sum(mapcols(f, x)), Zygote.Params([w,x]))
     @test grad_mapcols[x] ≈ gradx
-    @test_broken grad_mapcols[w] ≈ gradw # grad_mapcols[w] === nothing
+    @test grad_mapcols[w] == nothing # bug or a feature?
 
-    grad_slicemap = Zygote.gradient(() -> sum(slicemap(F(w), x, dims=1)), Zygote.Params([w,x]))
+    grad_slicemap = Zygote.gradient(() -> sum(slicemap(f, x, dims=1)), Zygote.Params([w,x]))
     @test grad_slicemap[x] ≈ gradx
-    @test_broken grad_slicemap[w] ≈ gradw # wrong numbers
+    @test grad_slicemap[w] ≈ gradw
     @test gradw ≈ Zygote.gradient(w -> sum(slicemap(F(w), x, dims=1)), w)[1]
+    # Using F(w) with Params() gives wrong answers:
+    # https://github.com/FluxML/Zygote.jl/issues/522#issuecomment-605935652
 
 end