Fix Zygote accum for CA in GPU broadcasting (#221)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "ComponentArrays"
 uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 authors = ["Jonnie Diegelman <[email protected]>"]
-version = "0.15.1"
+version = "0.15.2"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -21,6 +21,7 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 ComponentArraysAdaptExt = "Adapt"
@@ -30,6 +31,7 @@ ComponentArraysRecursiveArrayToolsExt = "RecursiveArrayTools"
 ComponentArraysReverseDiffExt = "ReverseDiff"
 ComponentArraysSciMLBaseExt = "SciMLBase"
 ComponentArraysTrackerExt = "Tracker"
+ComponentArraysZygoteExt = "Zygote"
 
 [compat]
 Adapt = "3"
@@ -46,6 +48,7 @@ SciMLBase = "1"
 StaticArraysCore = "1"
 StaticArrayInterface = "1"
 Tracker = "0.2"
+Zygote = "0.6"
 julia = "1.6"
 
 [extras]
@@ -58,3 +61,4 @@ ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

ext/ComponentArraysZygoteExt.jl

@@ -0,0 +1,13 @@
+module ComponentArraysZygoteExt
+
+using ComponentArrays, Zygote
+
+# For most cases this work. However, if the ComponentArray contains ROCArray, it fails to
+# compile the broadcast operation on AMDGPU. This will most likely be fixed with proper
+# broadcast mechanics in AMDGPU.jl but we can work around that in a harmless fashion for
+# now.
+function Zygote.accum(x::ComponentArray, ys::ComponentArray...)
+    return ComponentArray(Zygote.accum(getdata(x), getdata.(ys)...), getaxes(x))
+end
+
+end