Merge pull request #2 from mcabbott/zygoterules

ZygoteRules

Author: mcabbott

Project.toml

@@ -1,7 +1,7 @@
 name = "SliceMap"
 uuid = "82cb661a-3f19-5665-9e27-df437c7e54c8"
 authors = ["Michael Abbott"]
-version = "0.1.1"
+version = "0.1.2"
 
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -11,6 +11,7 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TensorCast = "02d47bb6-7ce6-556a-be16-bb1710789e2b"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 julia = "1"

src/SliceMap.jl

@@ -8,6 +8,9 @@ using MacroTools, Requires, TensorCast, JuliennedArrays
 using Tracker
 using Tracker: TrackedMatrix, track, @grad, data
 
+using ZygoteRules
+using ZygoteRules: pullback, @adjoint
+
 #========== Reverse, Eachslice ==========#
 
 """
@@ -36,6 +39,9 @@ _mapcols(map::Function, f::Function, M::TrackedMatrix, args...) = track(_mapcols
 @grad _mapcols(map::Function, f::Function, M::AbstractMatrix, args...) =
     ∇mapcols(map, map(col -> Tracker.forward(x -> surevec(f(x, args...)), col), eachcol(data(M))), args...)
 
+@adjoint _mapcols(map::Function, f::Function, M::AbstractMatrix, args...) =
+    ∇mapcols(map, map(col -> ZygoteRules.pullback(x -> surevec(f(x, args...)), col), eachcol(M)), args)
+
 function ∇mapcols(bigmap, forwards, args...)
     reduce(hcat, map(data∘first, forwards)), Δ -> begin
         cols = bigmap((fwd, Δcol) -> data(last(fwd)(Δcol)[1]), forwards, eachcol(data(Δ)))
@@ -56,6 +62,9 @@ maprows(f::Function, M::TrackedMatrix, args...) = track(maprows, f, M, args...)
 @grad maprows(f::Function, M::AbstractMatrix, args...) =
     ∇maprows(map(row -> Tracker.forward(x -> surevec(f(x, args...)), row), eachrow(data(M))), args)
 
+@adjoint maprows(f::Function, M::AbstractMatrix, args...) =
+    ∇maprows(map(row -> ZygoteRules.pullback(x -> surevec(f(x, args...)), row), eachrow(M)), args)
+
 function ∇maprows(forwards, args)
     reduce(vcat, map(transpose∘data∘first, forwards)), Δ -> begin
         rows = map((fwd, Δrow) -> data(last(fwd)(Δrow)[1]), forwards, eachrow(data(Δ)))
@@ -77,6 +86,7 @@ function slicemap(f::Function, A::AbstractArray{T,N}, args...; dims) where {T,N}
     C = [ f(slice, args...) for slice in B ]
     TensorCast.glue(C, code)
 end
+# TODO switch to JuliennedArrays, then rm TensorCast dep
 
 #========== Forward, Static ==========#
 
@@ -111,6 +121,9 @@ _MapCols(map::Function, f::Function, M::TrackedMatrix, dval, args...) =
 @grad _MapCols(map::Function, f::Function, M::TrackedMatrix, dval, args...) =
     ∇MapCols(map, f, M, dval, args...)
 
+@adjoint _MapCols(map::Function, f::Function, M::Matrix, dval, args...) =
+    ∇MapCols(map, f, M, dval, args...)
+
 function ∇MapCols(bigmap::Function, f::Function, M::AbstractMatrix{T}, dval::Val{d}, args...) where {T,d}
     d == size(M,1) || error("expected M with $d rows")
     k = size(M,2)
@@ -142,7 +155,34 @@ end
 
 # @require Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
-@init @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" include("zygote.jl")
+# @init @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" include("zygote.jl")
+# Now using ZygoteRules instead, mapcols etc above.
+
+#= TensorCast =#
+# These could move there, TODO
+
+@adjoint TensorCast.sliceview(A::AbstractArray, code::Tuple) =
+    TensorCast.sliceview(A, code), Δ -> (TensorCast.glue(Δ, code), nothing)
+
+@adjoint TensorCast.red_glue(A::AbstractArray, code::Tuple) =
+    TensorCast.red_glue(A, code), Δ -> (TensorCast.sliceview(Δ, code), nothing)
+
+@adjoint TensorCast.copy_glue(A::AbstractArray, code::Tuple) =
+    TensorCast.copy_glue(A, code), Δ -> (TensorCast.sliceview(Δ, code), nothing)
+
+#= JuliennedArrays =#
+
+@adjoint JuliennedArrays.Slices(whole, along...) =
+    Slices(whole, along...), Δ -> (Align(Δ, along...), map(_->nothing, along)...)
+
+@adjoint JuliennedArrays.Align(whole, along...) =
+    Align(whole, along...), Δ -> (Slices(Δ, along...), map(_->nothing, along)...)
+
+#= Base =#
+
+@adjoint Base.reduce(::typeof(hcat), V::AbstractVector{<:AbstractVector}) =
+    reduce(hcat, V), dV -> (nothing, collect(eachcol(dV)),)
+
 
 #========== Experimenting with gradients for for eachslice / reduce ==========#
 
@@ -159,6 +199,9 @@ gluecol(V::AbstractVector{<:TrackedVector}) = track(gluecol, V)
 end
 =#
 
+@adjoint gluecol(V::AbstractVector) =
+    gluecol(V), ΔM -> (collect(eachcol(ΔM)),) # does work!
+
 function mapcols2(f, A)
     cols = [A[:,c] for c=1:size(A,2)]
     res = f.(cols)
@@ -197,6 +240,14 @@ end
 
 collecteachcol(x) = collect(eachcol(x))
 
+@adjoint function collecteachcol(x)
+    collecteachcol(x), dy -> begin
+        dx = _zero(x) # _zero is not in ZygoteRules, TODO
+        foreach(copyto!, collecteachcol(dx), dy)
+        (dx,)
+    end
+end
+
 function mapcols6(f, A)
     cols = collecteachcol(A)
     res = map(f, cols)

src/zygote.jl

@@ -1,4 +1,7 @@
 
+# Note to self -- ZygoteRules doesn't have forward, only @adjoint, so it's no help.
+# Later: it now has pullback, so that should work? Once registered... all copied in.
+
 using .Zygote
 using .Zygote: @adjoint, _zero, forward