Merge pull request #679 from cossio/main

fix cat rrule

Author: mzgubic

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.44.6"
+version = "1.44.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/rulesets/Base/array.jl

@@ -347,7 +347,8 @@ end
 function rrule(::typeof(cat), Xs...; dims)
     Y = cat(Xs...; dims=dims)
     Base.require_one_based_indexing(Y)
-    cdims = dims isa Val ? Int(_val(dims)) : dims isa Integer ? Int(dims) : Tuple(dims)
+    _cdims = dims isa Val ? _val(dims) : dims
+    cdims = _cdims isa Integer ? Int(_cdims) : Tuple(_cdims)
     ndimsY = Val(ndims(Y))
     sizes = map(_catsize, Xs)
     project_Xs = map(ProjectTo, Xs)

test/rulesets/Base/array.jl

@@ -237,6 +237,10 @@ end
     @gpu test_rrule(cat, rand(2, 4), rand(2); fkwargs=(dims=Val(2),))
     test_rrule(cat, rand(), rand(2, 3); fkwargs=(dims=[1,2],))
     test_rrule(cat, rand(1), rand(3, 2, 1); fkwargs=(dims=(1,2),), check_inferred=false) # infers Tuple{Zero, Vector{Float64}, Any}
+    
+    if VERSION ≥ v"1.8" # Val(tuple) dims support was added in v1.8
+        test_rrule(cat, randn(3,2,4), randn(3,2,4); fkwargs=(dims=Val((1,2)),)) #678
+    end
 
     test_rrule(cat, rand(2, 2), rand(2, 2)'; fkwargs=(dims=1,))
     # inference on exotic array types