Widen cat rules (#614)

* widen cat rules

* cat+version

Author: mcabbott

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.29.0"
+version = "1.30.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Base/array.jl

@@ -209,10 +209,14 @@ function frule((_, ẋs...), ::typeof(hcat), xs...)
     return hcat(xs...), hcat(_instantiate_zeros(ẋs, xs)...)
 end
 
-function rrule(::typeof(hcat), Xs::Union{AbstractArray, Number}...)
+# All the [hv]cat functions treat anything that's not an array as a scalar. 
+_catsize(x) = ()
+_catsize(x::AbstractArray) = size(x)
+
+function rrule(::typeof(hcat), Xs...)
     Y = hcat(Xs...)  # note that Y always has 1-based indexing, even if X isa OffsetArray
     ndimsY = Val(ndims(Y))  # this avoids closing over Y, Val() is essential for type-stability
-    sizes = map(size, Xs)   # this avoids closing over Xs
+    sizes = map(_catsize, Xs)   # this avoids closing over Xs
     project_Xs = map(ProjectTo, Xs)
     function hcat_pullback(ȳ)
         dY = unthunk(ȳ)
@@ -279,10 +283,10 @@ function frule((_, ẋs...), ::typeof(vcat), xs...)
     return vcat(xs...), vcat(_instantiate_zeros(ẋs, xs)...)
 end
 
-function rrule(::typeof(vcat), Xs::Union{AbstractArray, Number}...)
+function rrule(::typeof(vcat), Xs...)
     Y = vcat(Xs...)
     ndimsY = Val(ndims(Y))
-    sizes = map(size, Xs)
+    sizes = map(_catsize, Xs)
     project_Xs = map(ProjectTo, Xs)
     function vcat_pullback(ȳ)
         dY = unthunk(ȳ)
@@ -342,11 +346,11 @@ function frule((_, ẋs...), ::typeof(cat), xs...; dims)
     return cat(xs...; dims), cat(_instantiate_zeros(ẋs, xs)...; dims)
 end
 
-function rrule(::typeof(cat), Xs::Union{AbstractArray, Number}...; dims)
+function rrule(::typeof(cat), Xs...; dims)
     Y = cat(Xs...; dims=dims)
     cdims = dims isa Val ? Int(_val(dims)) : dims isa Integer ? Int(dims) : Tuple(dims)
     ndimsY = Val(ndims(Y))
-    sizes = map(size, Xs)
+    sizes = map(_catsize, Xs)
     project_Xs = map(ProjectTo, Xs)
     function cat_pullback(ȳ)
         dY = unthunk(ȳ)
@@ -384,11 +388,11 @@ function frule((_, _, ẋs...), ::typeof(hvcat), rows, xs...)
     return hvcat(rows, xs...), hvcat(rows, _instantiate_zeros(ẋs, xs)...)
 end
 
-function rrule(::typeof(hvcat), rows, values::Union{AbstractArray, Number}...)
+function rrule(::typeof(hvcat), rows, values...)
     Y = hvcat(rows, values...)
     cols = size(Y,2)
     ndimsY = Val(ndims(Y))
-    sizes = map(size, values)
+    sizes = map(_catsize, values)
     project_Vs = map(ProjectTo, values)
     function hvcat_pullback(dY)
         prev = fill(0, 2)

test/rulesets/Base/array.jl

@@ -166,7 +166,8 @@ end
     test_rrule(hcat, rand(3,1,1,2), rand(3,3,1,2))
 
     # mix types
-    test_rrule(hcat, rand(2, 2), rand(2, 2)')
+    test_rrule(hcat, rand(1, 3), rand(2)')
+    test_rrule(hcat, rand(1), (nothing, rand()), check_inferred=false)
 end
 
 @testset "reduce hcat" begin
@@ -203,9 +204,10 @@ end
     test_rrule(vcat, rand(), rand(3), rand(3,1,1))
     test_rrule(vcat, rand(3,1,2), rand(4,1,2))
 
-
     # mix types
     test_rrule(vcat, rand(2, 2), rand(2, 2)')
+    test_rrule(vcat, rand(), rand() => rand(); check_inferred=false)
+    test_rrule(vcat, rand(3), (rand(), nothing), pi/2; check_inferred=false)
 end
 
 @testset "reduce vcat" begin
@@ -235,6 +237,7 @@ end
     test_rrule(cat, rand(2, 2), rand(2, 2)'; fkwargs=(dims=1,))
     # inference on exotic array types
     test_rrule(cat, @SArray(rand(3, 2, 1)), @SArray(rand(3, 2, 1)); fkwargs=(dims=Val(2),))
+    test_rrule(cat, pi/2, rand(1,3), (4.5,); fkwargs=(;dims=(2,)), check_inferred=false)
 end
 
 @testset "hvcat" begin
@@ -249,6 +252,7 @@ end
 
     # mix types (adjoint and transpose)
     test_rrule(hvcat, 1, rand(3)', transpose(rand(3)) ⊢ rand(1,3))
+    test_rrule(hvcat, (1,2), rand(2)', (3.4, 5.6), 7.8; check_inferred=false)
 end
 
 @testset "reverse" begin