remove some ::Function, and add tests f(w)(x)

Author: Michael Abbott

src/SliceMap.jl

@@ -26,23 +26,23 @@ 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.
 """
-mapcols(f::Function, M, args...) = _mapcols(map, f, M, args...)
-tmapcols(f::Function, M, args...) = _mapcols(threadmap, f, M, args...)
+mapcols(f, M, args...) = _mapcols(map, f, M, args...)
+tmapcols(f, M, args...) = _mapcols(threadmap, f, M, args...)
 
-function _mapcols(map::Function, f::Function, M::AbstractMatrix, args...)
+function _mapcols(map::Function, f, M::AbstractMatrix, args...)
     res = map(col -> _vec(f(col, args...)), eachcol(M))
     eltype(res) <: AbstractVector ? reduce(hcat, res) : reshape(res,1,:)
 end
 
 _vec(x) = x
 _vec(A::AbstractArray) = vec(A) # to allow f vector -> matrix, by reshaping
 
-_mapcols(map::Function, f::Function, M::TrackedMatrix, args...) = track(_mapcols, map, f, M, args...)
+_mapcols(map::Function, f, M::TrackedMatrix, args...) = track(_mapcols, map, f, M, args...)
 
-@grad _mapcols(map::Function, f::Function, M::AbstractMatrix, args...) =
+@grad _mapcols(map::Function, f, M::AbstractMatrix, args...) =
     ∇mapcols(map, map(col -> Tracker.forward(x -> _vec(f(x, args...)), col), eachcol(data(M))), args...)
 
-@adjoint _mapcols(map::Function, f::Function, M::AbstractMatrix, args...) =
+@adjoint _mapcols(map::Function, f, M::AbstractMatrix, args...) =
     ∇mapcols(map, map(col -> ZygoteRules.pullback(x -> _vec(f(x, args...)), col), eachcol(M)), args)
 
 function ∇mapcols(bigmap, forwards, args...)
@@ -91,7 +91,7 @@ e.g. if `dims=(2,4)` then `f` must map matrices to matrices.
 
 The gradient is for Zygote only.
 """
-function slicemap(f::Function, A::AbstractArray{T,N}, args...; dims) where {T,N}
+function slicemap(f, A::AbstractArray{T,N}, args...; dims) where {T,N}
     code = ntuple(d -> d in dims ? True() : False(), N)
     B = JuliennedArrays.Slices(A, code...)
     C = [ f(slice, args...) for slice in B ]

test/runtests.jl

@@ -133,3 +133,27 @@ end
     @test grad ≈ Zygote.gradient(m -> sum(sin, j3(fun, m)), ten)[1]
 
 end
+@testset "gradient of the function" begin
+
+    struct F W end
+    (f::F)(x) = f.W * x # toy version of e.g. Flux.Dense
+    w = rand(3,2)
+    x = rand(2,5)
+    gradx = ForwardDiff.gradient(x -> sum(mapslices(F(w), x, dims=1)), x)
+    gradw = ForwardDiff.gradient(w -> sum(mapslices(F(w), x, dims=1)), w)
+
+    wp = Tracker.param(w)
+    xp = Tracker.param(x)
+    Tracker.back!(sum(mapcols(F(wp), xp)))
+    @test Tracker.grad(xp) ≈ gradx
+    @test_broken Tracker.grad(wp) ≈ gradw # zero
+
+    grad_mapcols = Zygote.gradient(() -> sum(mapcols(F(w), x)), Zygote.Params([w,x]))
+    @test grad_mapcols[x] ≈ gradx
+    @test_broken grad_mapcols[w] ≈ gradw # grad_mapcols[w] === nothing
+
+    grad_slicemap = Zygote.gradient(() -> sum(slicemap(F(w), x, dims=1)), Zygote.Params([w,x]))
+    @test grad_slicemap[x] ≈ gradx
+    @test_broken grad_slicemap[w] ≈ gradw # wrong numbers
+
+end