Merge pull request #7 from mcabbott/impure

Allow f::Any, possibly containing parameters

Author: mcabbott

.travis.yml

@@ -1,11 +1,9 @@
 language: julia
 os:
   - linux
-  - osx
-  - windows
 julia:
   - 1.1
-  - 1.2
+  - 1.4
   - nightly
 
 matrix:

Project.toml

@@ -1,7 +1,7 @@
 name = "SliceMap"
 uuid = "82cb661a-3f19-5665-9e27-df437c7e54c8"
 authors = ["Michael Abbott"]
-version = "0.2.1"
+version = "0.2.2"
 
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

src/SliceMap.jl

@@ -25,24 +25,26 @@ 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::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...)
@@ -90,8 +92,11 @@ 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::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,37 @@ 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 Tracker.grad(wp) == 0 .* gradw # bug or a feature?
+
+    # 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 grad_mapcols[w] == nothing # bug or a feature?
+
+    grad_slicemap = Zygote.gradient(() -> sum(slicemap(f, x, dims=1)), Zygote.Params([w,x]))
+    @test grad_slicemap[x] ≈ gradx
+    @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