rename with unzip

Author: mcabbott

src/ChainRules.jl

@@ -24,7 +24,7 @@ using ChainRulesCore: derivatives_given_output
 const CommutativeMulNumber = Union{Real,Complex}
 
 # StructArrays
-include("tuplecast.jl")
+include("unzipped.jl")
 
 include("rulesets/Core/core.jl")
 

src/rulesets/Base/broadcast.jl

@@ -72,11 +72,11 @@ function split_bc_derivatives(f::F, args::Vararg{Any,N}) where {F,N}
     @debug("split broadcasting derivatives", f, N)
     ys = f.(args...)
     function bc_many_back(dys)
-        deltas = tuplecast(unthunk(dys), ys, args...) do dy, y, as...
+        deltas = unzip_broadcast(unthunk(dys), ys, args...) do dy, y, as...
             das = only(derivatives_given_output(y, f, as...))
             map(da -> dy * conj(da), das)  # possibly this * should be made nan-safe.
         end
-        dargs = map(unbroadcast, args, deltas)  # ideally sum in unbroadcast could be part of tuplecast?
+        dargs = map(unbroadcast, args, deltas)  # ideally sum in unbroadcast could be part of unzip_broadcast?
         return (TRI_NO..., dargs...)
     end
     bc_many_back(z::AbstractZero) = (TRI_NO..., map(Returns(z), args)...)
@@ -105,7 +105,7 @@ split_bc_forwards(cfg::RuleConfig{>:HasForwardsMode}, f::F, arg) where {F} = spl
 split_bc_forwards(cfg::RuleConfig, f::F, arg) where {F} = split_bc_inner(frule, cfg, f, arg)
 function split_bc_inner(frule_fun::R, cfg::RuleConfig, f::F, arg) where {R,F}
     @debug("split broadcasting forwards", frule_fun, f)
-    ys, ydots = tuplecast(arg) do a
+    ys, ydots = unzip_broadcast(arg) do a
         frule_fun(cfg, (NoTangent(), one(a)), f, a)
     end
     function back_forwards(dys)
@@ -124,11 +124,11 @@ end
 
 function split_bc_pullbacks(cfg::RCR, f::F, args::Vararg{Any,N}) where {F,N}
     @debug("split broadcasting generic", f, N)
-    ys3, backs = tuplecast(args...) do a...
+    ys3, backs = unzip_broadcast(args...) do a...
         rrule_via_ad(cfg, f, a...)
     end
     function back_generic(dys)
-        deltas = tuplecast(backs, unthunk(dys)) do back, dy  # (could be map, sizes match)
+        deltas = unzip_broadcast(backs, unthunk(dys)) do back, dy  # (could be map, sizes match)
             map(unthunk, back(dy))
         end
         dargs = map(unbroadcast, args, Base.tail(deltas))

src/tuplecast.jl renamed to src/unzipped.jl

@@ -1,29 +1,29 @@
 
 """
-    tuplecast(f, args...)
+    unzip_broadcast(f, args...)
 
 For a function `f` which returns a tuple, this is `== unzip(broadcast(f, args...))`,
 but performed using `StructArrays` for efficiency. Used in the gradient of broadcasting.
 
 # Examples
 ```
-julia> using ChainRules: tuplecast, unzip
+julia> using ChainRules: unzip_broadcast, unzip
 
-julia> tuplecast(x -> (x,2x), 1:3)
+julia> unzip_broadcast(x -> (x,2x), 1:3)
 ([1, 2, 3], [2, 4, 6])
 
-julia> mats = @btime tuplecast((x,y) -> (x+y, x-y), 1:1000, transpose(1:1000));  # 2 arrays, each 7.63 MiB
+julia> mats = @btime unzip_broadcast((x,y) -> (x+y, x-y), 1:1000, transpose(1:1000));  # 2 arrays, each 7.63 MiB
   min 1.776 ms, mean 20.421 ms (4 allocations, 15.26 MiB)
 
 julia> mats == @btime unzip(broadcast((x,y) -> (x+y, x-y), 1:1000, transpose(1:1000)))  # intermediate matrix of tuples
   min 2.660 ms, mean 40.007 ms (6 allocations, 30.52 MiB)
 true
 ```
 """
-function tuplecast(f::F, args...) where {F}
+function unzip_broadcast(f::F, args...) where {F}
     T = Broadcast.combine_eltypes(f, args)
     if isconcretetype(T)
-        T <: Tuple || throw(ArgumentError("""tuplecast(f, args) only works on functions returning a tuple,
+        T <: Tuple || throw(ArgumentError("""unzip_broadcast(f, args) only works on functions returning a tuple,
             but f = $(sprint(show, f)) returns type T = $T"""))
     end
     # TODO allow GPU arrays, possibly just as a fallback unzip, but see also: 
@@ -39,7 +39,7 @@ function tuplecast(f::F, args...) where {F}
     end
 end
 
-function ChainRulesCore.rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(tuplecast), f::F, args...) where {F}
+function ChainRulesCore.rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(unzip_broadcast), f::F, args...) where {F}
     y, back = rrule_via_ad(cfg, broadcast, f, args...)
     z = unzip(y)
     function untuplecast(dz)
@@ -53,23 +53,25 @@ function ChainRulesCore.rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(tuplec
 end
 
 # This is for testing, but the tests using it don't work.
-function rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(collect∘tuplecast), f, args...)
-    y, back = rrule(cfg, tuplecast, f, args...)
+function rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(collect∘unzip_broadcast), f, args...)
+    y, back = rrule(cfg, unzip_broadcast, f, args...)
     return collect(y), back
 end
 
+#=
+
 """
-    tuplemap(f, args...)
+    unzip_map(f, args...)
 
 For a function `f` which returns a tuple, this is `== unzip(map(f, args...))`,
 but performed using `StructArrays` for efficiency.
 
-Not in use at present, but see `tuplecast`.
+Not in use at present, but see `unzip_broadcast`.
 """
-function tuplemap(f::F, args...) where {F}
+function unzip_map(f::F, args...) where {F}
     T = Broadcast.combine_eltypes(f, args)
     if isconcretetype(T)
-        T <: Tuple || throw(ArgumentError("""tuplemap(f, args) only works on functions returning a tuple,
+        T <: Tuple || throw(ArgumentError("""unzip_map(f, args) only works on functions returning a tuple,
             but f = $(sprint(show, f)) returns type T = $T"""))
     end
     # if any(a -> a isa CuArray, args)
@@ -78,17 +80,19 @@ function tuplemap(f::F, args...) where {F}
     return StructArrays.components(StructArray(Iterators.map(f, args...)))
 end
 
-function ChainRulesCore.rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(tuplemap), f::F, xs...) where {F}
+function ChainRulesCore.rrule(cfg::RuleConfig{>:HasReverseMode}, ::typeof(unzip_map), f::F, xs...) where {F}
     y, back = rrule_via_ad(cfg, map, f, xs...)
     z = unzip(y)
-    function untuplemap(dz)
+    function ununzip_map(dz)
         # dy = StructArray(map(unthunk, dz))  # fails for e.g. StructArray(([1,2,3], ZeroTangent()))
         dy = broadcast(tuple, map(unthunk, dz)...)
         return back(dy)
     end
-    return z, untuplemap
+    return z, ununzip_map
 end
 
+=#
+
 """
     unzip(A)
 
@@ -114,7 +118,7 @@ function unzip(xs::AbstractArray)
     x1 = first(xs)
     x1 isa Tuple || throw(ArgumentError("unzip only accepts arrays of tuples"))
     N = length(x1)
-    return unzip(xs, Val(N))  # like Zygote's unzip, here this is the fallback case.
+    return unzip(xs, Val(N))  # like Zygote's unzip. Here this is the fallback case.
 end
 
 @generated function unzip(xs, ::Val{N}) where {N}

test/runtests.jl

@@ -59,7 +59,7 @@ end
     include_test("rulesets/Base/sort.jl")
     include_test("rulesets/Base/broadcast.jl")
 
-    include_test("tuplecast.jl")  # used primarily for broadcast
+    include_test("unzipped.jl")  # used primarily for broadcast
 
     println()
 

test/tuplecast.jl renamed to test/unzipped.jl

@@ -1,8 +1,8 @@
 
-using ChainRules: tuplecast, unzip, tuplemap
+using ChainRules: unzip_broadcast, unzip #, unzip_map
 
-@testset "tuplecast.jl" begin
-    @testset "basics: $(sprint(show, fun))" for fun in [tuplemap, tuplecast, unzip∘map, unzip∘broadcast]
+@testset "unzip_broadcast.jl" begin
+    @testset "basics: $(sprint(show, fun))" for fun in [unzip_broadcast, unzip∘map, unzip∘broadcast] # unzip_map, 
         @test_throws Exception fun(sqrt, 1:3)
 
         @test fun(tuple, 1:3, 4:6) == ([1, 2, 3], [4, 5, 6])
@@ -17,7 +17,7 @@ using ChainRules: tuplecast, unzip, tuplemap
             @test fun(tuple, [1,2,3], [4 5]) == ([1 1; 2 2; 3 3], [4 5; 4 5; 4 5])
         end
 
-        if fun == tuplemap
+        if fun == unzip_map
             @test_broken fun(tuple, (1,2,3), (4,5,6)) == ((1, 2, 3), (4, 5, 6))
         elseif fun == unzip∘map
             @test fun(tuple, (1,2,3), (4,5,6)) == ((1, 2, 3), (4, 5, 6))
@@ -32,19 +32,19 @@ using ChainRules: tuplecast, unzip, tuplemap
     @testset "rrules" begin
         # These exist to allow for second derivatives
 
-        # test_rrule(collect∘tuplecast, tuple, [1,2,3.], [4,5,6.], collectheck_inferred=false) # return type Tuple{NoTangent, NoTangent, Vector{Float64}, Vector{Float64}} does not match inferred return type NTuple{4, Any}
+        # test_rrule(collect∘unzip_broadcast, tuple, [1,2,3.], [4,5,6.], collectheck_inferred=false) # return type Tuple{NoTangent, NoTangent, Vector{Float64}, Vector{Float64}} does not match inferred return type NTuple{4, Any}
 
-        y1, bk1 = rrule(CFG, tuplecast, tuple, [1,2,3.0], [4,5,6.0])
+        y1, bk1 = rrule(CFG, unzip_broadcast, tuple, [1,2,3.0], [4,5,6.0])
         @test y1 == ([1, 2, 3], [4, 5, 6])
         @test bk1(([1,10,100.0], [7,8,9.0]))[3] ≈ [1,10,100]
 
         # bk1(([1,10,100.0], NoTangent()))  # DimensionMismatch in FiniteDifferences
 
-        y2, bk2 = rrule(CFG, tuplecast, tuple, [1,2,3.0], [4 5.0], 6.0)
+        y2, bk2 = rrule(CFG, unzip_broadcast, tuple, [1,2,3.0], [4 5.0], 6.0)
         @test y2 == ([1 1; 2 2; 3 3], [4 5; 4 5; 4 5], [6 6; 6 6; 6 6])
         @test bk2(y2)[5] ≈ 36
 
-        y4, bk4 = rrule(CFG, tuplemap, tuple, [1,2,3.0], [4,5,6.0])
+        y4, bk4 = rrule(CFG, unzip_map, tuple, [1,2,3.0], [4,5,6.0])
         @test y4 == ([1, 2, 3], [4, 5, 6])
         @test bk4(([1,10,100.0], [7,8,9.0]))[3] ≈ [1,10,100]
     end