add test for ComponentArrays

Author: CarloLucibello

Manifest.toml

@@ -95,6 +95,12 @@ version = "3.30.0"
 deps = ["Artifacts", "Libdl"]
 uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
 
+[[ComponentArrays]]
+deps = ["ArrayInterface", "LinearAlgebra", "Requires"]
+git-tree-sha1 = "76495e7a7e47abc3771d70c782d5f6e66f114d36"
+uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+version = "0.10.5"
+
 [[DataAPI]]
 git-tree-sha1 = "dfb3b7e89e395be1e25c2ad6d7690dc29cc53b1d"
 uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"

Project.toml

@@ -46,11 +46,12 @@ Zygote = "0.6"
 julia = "1.6"
 
 [extras]
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Documenter", "IterTools", "LinearAlgebra", "FillArrays"]
+test = ["Test", "Documenter", "IterTools", "LinearAlgebra", "FillArrays", "ComponentArrays"]

test/optimise.jl

@@ -1,7 +1,7 @@
 using Flux.Optimise
 using Flux.Optimise: runall
 using Flux: Params, gradient
-import FillArrays
+import FillArrays, ComponentArrays
 using Test
 using Random
 
@@ -143,7 +143,7 @@ end
   @test norm(w̄_norm) <= 1
 end
 
-@testset "handle Fills from Zygote" begin
+@testset "update!: handle Fills from Zygote" begin
   w = randn(10,10)
   wold = copy(w)
   g = FillArrays.Ones(size(w))
@@ -169,3 +169,16 @@ end
   Flux.update!(opt, θ, gs)
   @test w ≈ wold .- 0.1 
 end
+
+@testset "update!: handle ComponentArrays" begin
+  w = ComponentArrays.ComponentArray(a=1.0, b=[2, 1, 4], c=(a=2, b=[1, 2]))
+  wold = deepcopy(w)
+  θ = Flux.params([w])
+  gs = gradient(() -> sum(w.a) + sum(w.c.b), θ)
+  opt = Descent(0.1)
+  Flux.update!(opt, θ, gs)
+  @test w.a ≈ wold.a .- 0.1
+  @test w.b ≈ wold.b
+  @test w.c.b ≈ wold.c.b .- 0.1
+  @test w.c.a ≈ wold.c.a
+end