+ JuliennedArrays

Author: Michael Abbott

Project.toml

@@ -5,6 +5,7 @@ version = "0.1.0"
 
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JuliennedArrays = "5cadff95-7770-533d-a838-a1bf817ee6e0"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

README.md

@@ -14,7 +14,7 @@ maprows(f, M) ≈ mapreduce(f, vcat, eachrow(M))
 slicemap(f, A; dims) ≈ mapslices(f, A, dims)
 ```
 
-### An example
+### Simple example
 
 ```julia
 mat = rand(1:9, 3,10)
@@ -56,27 +56,44 @@ mat1k = rand(3,1000);
 @btime Zygote.gradient(m -> sum(sin, MapCols{3}(fun, m)), $mat1k);             #     245.550 μs
 ```
 
-It also provides Zygote gradients for the slice/glue functions in 
+### Other packages
+
+This package also provides Zygote gradients for the slice/glue functions in 
 [TensorCast](https://github.com/mcabbott/TensorCast.jl),
 which can be used to write many mapslices-like operations.
 (The function `slicemap(f, A, dims)` uses these functions, without having to write index notation.)
 
 ```julia
 using TensorCast
-@cast [i,j] := fun(mat[:,j])[i]                       # same as mapcols
+@cast [i,j] := fun(mat[:,j])[i]                        # same as mapcols
 
 tcm(mat) = @cast out[i,j] := fun(mat[:,j])[i]
 Zygote.gradient(m -> sum(sin, tcm(m)), mat)[1]
 
-@btime tcm($mat1k)                                    # 407.176 μs
-@btime Zygote.gradient(m -> sum(sin, tcm(m)), $mat1k) # 19.086 ms
+@btime tcm($mat1k)                                     #    407.176 μs
+@btime Zygote.gradient(m -> sum(sin, tcm(m)), $mat1k); # 19.086 ms
+```
+
+Similar gradients work for the Slice/Align functions in 
+[JuliennedArrays](https://github.com/bramtayl/JuliennedArrays.jl),
+so it defines these too:
+
+```julia
+using JuliennedArrays
+jumap(f,m) = Align(map(f, Slices(m, True(), False())), True(), False())
+jumap(fun, mat)                                               # same as mapcols
+Zygote.gradient(m -> sum(sin, jumap(fun, m)), mat)[1]
+
+@btime jumap(fun, $mat1k);                                    #    408.259 μs
+@btime Zygote.gradient(m -> sum(sin, jumap(fun, m)), $mat1k); # 18.638 ms
 ```
 
 ### Elsewhere
 
-Issues about mapslices:
+About mapslices:
 * https://github.com/FluxML/Zygote.jl/issues/92
 * https://github.com/FluxML/Flux.jl/issues/741
+* https://github.com/JuliaLang/julia/issues/29146
 
 Other packages which define gradients of possible interest:
 * https://github.com/GiggleLiu/LinalgBackwards.jl

src/SliceMap.jl

@@ -3,8 +3,9 @@ module SliceMap
 
 export mapcols, MapCols, maprows, slicemap
 
-using MacroTools, Requires, WeightedArrays, TensorCast, Tracker
+using MacroTools, Requires, WeightedArrays, TensorCast, JuliennedArrays
 
+using Tracker
 using Tracker: TrackedMatrix, track, @grad, data
 
 #========== Reverse, Eachslice ==========#

src/zygote.jl

@@ -14,12 +14,16 @@ using .Zygote: @adjoint, _zero, forward
 
 #===== TensorCast =====#
 
+# @adjoint function TensorCast.sliceview(A::AbstractArray, code::Tuple)
+#     TensorCast.sliceview(A, code), Δ -> begin
+#         dA = _zero(A)
+#         foreach(copyto!, TensorCast.sliceview(dA, code), Δ)
+#         (dA, nothing)
+#     end
+# end
+
 @adjoint function TensorCast.sliceview(A::AbstractArray, code::Tuple)
-    TensorCast.sliceview(A, code), Δ -> begin
-        dA = _zero(A)
-        foreach(copyto!, TensorCast.sliceview(dA, code), Δ)
-        (dA, nothing)
-    end
+    TensorCast.sliceview(A, code), Δ -> (TensorCast.glue(Δ, code), nothing)
 end
 
 @adjoint function TensorCast.red_glue(A::AbstractArray, code::Tuple)
@@ -30,6 +34,16 @@ end
     TensorCast.copy_glue(A, code), Δ -> (TensorCast.sliceview(Δ, code), nothing)
 end
 
+#===== JuliennedArrays =====#
+
+@adjoint function Slices(whole, along...)
+    Slices(whole, along...), Δ -> (Align(Δ, along...), map(_->nothing, along)...)
+end
+
+@adjoint function Align(whole, along...)
+    Align(whole, along...), Δ -> (Slices(Δ, along...), map(_->nothing, along)...)
+end
+
 #===== Misc Base =====#
 
 @adjoint function Base.reduce(::typeof(hcat), V::AbstractVector{<:AbstractVector})

test/runtests.jl

@@ -1,7 +1,7 @@
 
 using SliceMap
 using Test
-using ForwardDiff, Tracker, Zygote, TensorCast
+using ForwardDiff, Tracker, Zygote, TensorCast, JuliennedArrays
 
 Zygote.refresh()
 
@@ -26,8 +26,12 @@ Zygote.refresh()
     @test res ≈ tcm(mat)
     @test grad ≈ Zygote.gradient(m -> sum(sin, tcm(m)), mat)[1]
 
+    jcols(f,m) = Align(map(f, Slices(m, True(), False())), True(), False())
+    @test res ≈ jcols(fun, mat)
+    @test grad ≈ Zygote.gradient(m -> sum(sin, jcols(fun, m)), mat)[1]
+
 end
-@testset "columns, scalar" begin
+@testset "columns -> scalar" begin
 
     mat = rand(1:9, 3,10)
     fun(x) = sum(x) # different function!
@@ -49,7 +53,7 @@ end
     # @test grad ≈ Zygote.gradient(m -> sum(sin, tcm3(m)), mat)[1]
 
 end
-@testset "columns, matrix" begin
+@testset "columns -> matrix" begin
 
     mat = rand(1:9, 3,10)
     fun(x) = x .* x' # different function! vector -> matrix
@@ -87,16 +91,27 @@ end
     # @test res ≈ tcm2(mat)
     # @test grad ≈ Zygote.gradient(m -> sum(sin, tcm2(m)), mat)[1]
 
+    jrows(f,m) = Align(map(f, Slices(m, False(), True())), False(), True())
+    @test res ≈ jrows(fun, mat)
+    @test grad ≈ Zygote.gradient(m -> sum(sin, jrows(fun, m)), mat)[1]
+
+
 end
-@testset "slices" begin
+@testset "slices of a 4-tensor" begin
 
     ten = randn(3,4,5,2)
-    fun(x) = sqrt(3) .+ x.^3 ./ (sum(x)^2)
+    fun(x::AbstractVector) = sqrt(3) .+ x.^3 ./ (sum(x)^2)
     res = mapslices(fun, ten, dims=3)
 
     @test res ≈ slicemap(fun, ten, dims=3)
 
     grad = ForwardDiff.gradient(x -> sum(sin, slicemap(fun, x, dims=3)), ten)
     @test grad ≈ Zygote.gradient(x -> sum(sin, slicemap(fun, x, dims=3)), ten)[1]
 
+    jthree(f,m) = Align(map(f,
+        Slices(m, False(), False(), True(), False())
+        ), False(), False(), True(), False())
+    @test res ≈ jthree(fun, ten)
+    @test grad ≈ Zygote.gradient(m -> sum(sin, jthree(fun, m)), ten)[1]
+
 end