args bug + tests

Author: Michael Abbott

src/SliceMap.jl

@@ -37,9 +37,9 @@ surevec(A) = vec(A)      # to allow f vector -> matrix, by reshaping
 _mapcols(map::Function, f::Function, M::TrackedMatrix, args...) = track(_mapcols, map, f, M, args...)
 
 @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)
+    ∇mapcols(map, map(col -> Tracker.forward(x -> surevec(f(x, args...)), col), eachcol(data(M))), args...)
 
-function ∇mapcols(bigmap, forwards, 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(Δ)))
         (nothing, nothing, reduce(hcat, cols), map(_->nothing, args)...)
@@ -220,12 +220,15 @@ end
 # What KissThreading does is much more complicated, perhaps worth investigating:
 # https://github.com/mohamed82008/KissThreading.jl/blob/master/src/KissThreading.jl
 
+# BTW I do the first one because some diffeq maps are infer to ::Any
+# else you could use Core.Compiler.return_type(f, Tuple{eltype(x)})
+
 """
     threadmap(f, A)
     threadmap(f, A, B)
 
 Simple version of `map` using a `Threads.@threads` loop;
-only for vectors & only two of them, of nonzero length,
+only for vectors & really at most two of them, of nonzero length,
 with all outputs having the same type.
 """
 function threadmap(f::Function, vw::AbstractVector...)

test/runtests.jl

@@ -83,10 +83,32 @@ end
     @test grad ≈ Zygote.gradient(m -> sum(sin, mapcols(fun, m)), mat)[1]
     @test grad ≈ Zygote.gradient(m -> sum(sin, MapCols{3}(fun, m)), mat)[1]
 
-    # tcm4(mat) = @cast out[i⊗i′,j] := fun(mat[:,j])[i,i′] i:3, i′:3 # changed here too
+    # tcm4(mat) = @cast out[i⊗i′,j] := fun(mat[:,j])[i,i′]  i:3
     # @test res ≈ tcm4(mat)
     # @test grad ≈ Zygote.gradient(m -> sum(sin, tcm4(m)), mat)[1]
 
+end
+@testset "columns w args" begin
+
+    mat = randn(Float32, 3,10)
+    fun(x, s) = 1 .+ x .* s
+    res = mapslices(x -> vec(fun(x,5)), mat, dims=1)
+
+    @test res ≈ mapcols(fun, mat, 5)
+    @test res ≈ MapCols{3}(fun, mat, 5)
+
+    grad = ForwardDiff.gradient(m -> sum(sin, mapslices(x -> vec(fun(x,5)), m, dims=1)), mat)
+
+    @test grad ≈ Tracker.gradient(m -> sum(sin, mapcols(fun, m, 5)), mat)[1]
+    @test grad ≈ Tracker.gradient(m -> sum(sin, MapCols{3}(fun, m, 5)), mat)[1]
+
+    @test grad ≈ Zygote.gradient(m -> sum(sin, mapcols(fun, m, 5)), mat)[1]
+    @test grad ≈ Zygote.gradient(m -> sum(sin, MapCols{3}(fun, m, 5)), mat)[1]
+
+    # tcm5(mat) = @cast out[i,j] := fun(mat[:,j], 5)[i]
+    # @test res ≈ tcm5(mat)
+    # @test grad ≈ Zygote.gradient(m -> sum(sin, tcm5(m)), mat)[1]
+
 end
 @testset "rows" begin
 
@@ -108,7 +130,6 @@ end
     @test res ≈ jrows(fun, mat)
     @test grad ≈ Zygote.gradient(m -> sum(sin, jrows(fun, m)), mat)[1]
 
-
 end
 @testset "slices of a 4-tensor" begin