From 9e8a0b459f46e5e39bda9f36d6103e17a236da50 Mon Sep 17 00:00:00 2001
From: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
Date: Fri, 22 Apr 2022 17:19:09 -0400
Subject: [PATCH 1/2] test with GPUArrays

---
 Project.toml      |   5 ++-
 test/gpuarrays.jl | 103 ++++++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl  |   3 ++
 3 files changed, 110 insertions(+), 1 deletion(-)
 create mode 100644 test/gpuarrays.jl

diff --git a/Project.toml b/Project.toml
index c60e8de5..c7df4944 100644
--- a/Project.toml
+++ b/Project.toml
@@ -17,9 +17,12 @@ Zygote = "0.6.40"
 julia = "1.6"
 
 [extras]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "StaticArrays", "Zygote"]
+test = ["Adapt", "CUDA", "GPUArrays", "StaticArrays", "Test", "Zygote"]
diff --git a/test/gpuarrays.jl b/test/gpuarrays.jl
new file mode 100644
index 00000000..5aa9226f
--- /dev/null
+++ b/test/gpuarrays.jl
@@ -0,0 +1,103 @@
+using Optimisers
+using ChainRulesCore #, Functors, StaticArrays, Zygote
+using LinearAlgebra, Statistics, Test
+
+import CUDA
+if CUDA.functional()
+  using CUDA  # exports CuArray, etc
+  @info "starting CUDA tests"
+else
+  @info "CUDA not functional, testing via GPUArrays"
+  using GPUArrays
+  GPUArrays.allowscalar(false)
+
+  # GPUArrays provides a fake GPU array, for testing
+  jl_file = normpath(joinpath(pathof(GPUArrays), "..", "..", "test", "jlarray.jl"))
+  using Random, Adapt  # loaded within jl_file
+  include(jl_file)
+  using .JLArrays
+  cu = jl
+  CuArray{T,N} = JLArray{T,N}
+end
+
+@test cu(rand(3)) .+ 1 isa CuArray
+
+@testset "very basics" begin
+  m = (cu([1.0, 2.0]),)
+  mid = objectid(m[1])
+  g = (cu([25, 33]),)
+  o = Descent(0.1f0)
+  s = Optimisers.setup(o, m)
+  
+  s2, m2 = Optimisers.update(s, m, g)
+  @test Array(m[1]) == 1:2  # not mutated
+  @test m2[1] isa CuArray
+  @test Array(m2[1]) ≈ [1,2] .- 0.1 .* [25, 33]  atol=1e-6
+
+  s3, m3 = Optimisers.update!(s, m, g)
+  @test objectid(m3[1]) == mid
+  @test Array(m3[1]) ≈ [1,2] .- 0.1 .* [25, 33]  atol=1e-6
+
+  g4 = Tangent{typeof(m)}(g...)
+  s4, m4 = Optimisers.update!(s, (cu([1.0, 2.0]),), g4)
+  @test Array(m4[1]) ≈ [1,2] .- 0.1 .* [25, 33]  atol=1e-6
+end
+
+@testset "basic mixed" begin
+  # Works trivially as every element of the tree is either here or there
+  m = (device = cu([1.0, 2.0]), host = [3.0, 4.0], neither = (5, 6, sin))
+  s = Optimisers.setup(ADAM(0.1), m)
+  @test s.device.state[1] isa CuArray
+  @test s.host.state[1] isa Array
+
+  g = (device = cu([1, 0.1]), host = [1, 10], neither = nothing)
+  s2, m2 = Optimisers.update(s, m, g)
+  
+  @test m2.device isa CuArray
+  @test Array(m2.device) ≈ [0.9, 1.9]  atol=1e-6
+  
+  @test m2.host isa Array
+  @test m2.host ≈ [2.9, 3.9]
+end
+
+RULES = [
+  # Just a selection:
+  Descent(), ADAM(), RMSProp(), NADAM(),
+  # A few chained combinations:
+  OptimiserChain(WeightDecay(), ADAM(0.001)),
+  OptimiserChain(ClipNorm(), ADAM(0.001)),
+  OptimiserChain(ClipGrad(0.5), Momentum()),
+]
+
+name(o) = typeof(o).name.name  # just for printing testset headings
+name(o::OptimiserChain) = join(name.(o.opts), " → ")
+
+@testset "rules: simple sum" begin
+  @testset "$(name(o))" for o in RULES
+    m = cu(shuffle!(reshape(1:64, 8, 8) .+ 0.0))
+    s = Optimisers.setup(o, m)
+    for _ in 1:10
+      g = Zygote.gradient(x -> sum(abs2, x + x'), m)[1]
+      s, m = Optimisers.update!(s, m, g)
+    end
+    @test sum(m) < sum(1:64)
+  end
+end
+
+@testset "destructure GPU" begin
+  m = (x = cu(Float32[1,2,3]), y = (0, 99), z = cu(Float32[4,5]))
+  v, re = destructure(m)
+  @test v isa CuArray
+  @test re(2v).x isa CuArray
+end
+
+@testset "destructure mixed" begin
+  # Not sure what should happen here!
+  m_c1 = (x = cu(Float32[1,2,3]), y = Float32[4,5])
+  v, re = destructure(m_c1)
+  @test re(2v).x isa CuArray
+  @test_broken re(2v).y isa Array
+  
+  m_c2 = (x = Float32[1,2,3], y = cu(Float32[4,5]))
+  @test_skip destructure(m_c2)  # ERROR: Scalar indexing
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 77c8af7a..8d9cdca2 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -216,4 +216,7 @@ Optimisers.trainable(x::TwoThirds) = (a = x.a,)
   @testset verbose=true "Optimisation Rules" begin
     include("rules.jl")
   end
+  @testset verbose=true "GPU" begin
+    include("gpuarrays.jl")
+  end
 end

From 307beef51418ab6990d63a5aa34087800511c600 Mon Sep 17 00:00:00 2001
From: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
Date: Sun, 24 Jul 2022 15:02:43 -0400
Subject: [PATCH 2/2] use JLArrays, add gradient test

---
 Project.toml      |  5 ++---
 test/gpuarrays.jl | 23 +++++++++++------------
 2 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/Project.toml b/Project.toml
index c7df4944..a6770ef1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -17,12 +17,11 @@ Zygote = "0.6.40"
 julia = "1.6"
 
 [extras]
-Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Adapt", "CUDA", "GPUArrays", "StaticArrays", "Test", "Zygote"]
+test = ["CUDA", "JLArrays", "StaticArrays", "Test", "Zygote"]
diff --git a/test/gpuarrays.jl b/test/gpuarrays.jl
index 5aa9226f..ff1724d7 100644
--- a/test/gpuarrays.jl
+++ b/test/gpuarrays.jl
@@ -1,25 +1,19 @@
 using Optimisers
-using ChainRulesCore #, Functors, StaticArrays, Zygote
-using LinearAlgebra, Statistics, Test
+using ChainRulesCore, Zygote
+using Test
 
 import CUDA
 if CUDA.functional()
   using CUDA  # exports CuArray, etc
-  @info "starting CUDA tests"
+  CUDA.allowscalar(false)
 else
-  @info "CUDA not functional, testing via GPUArrays"
-  using GPUArrays
-  GPUArrays.allowscalar(false)
+  @info "CUDA not functional, testing with JLArrays instead"
+  using JLArrays
+  JLArrays.allowscalar(false)
 
-  # GPUArrays provides a fake GPU array, for testing
-  jl_file = normpath(joinpath(pathof(GPUArrays), "..", "..", "test", "jlarray.jl"))
-  using Random, Adapt  # loaded within jl_file
-  include(jl_file)
-  using .JLArrays
   cu = jl
   CuArray{T,N} = JLArray{T,N}
 end
-
 @test cu(rand(3)) .+ 1 isa CuArray
 
 @testset "very basics" begin
@@ -89,6 +83,11 @@ end
   v, re = destructure(m)
   @test v isa CuArray
   @test re(2v).x isa CuArray
+
+  dm = gradient(m -> sum(abs2, destructure(m)[1]), m)[1]
+  @test dm.z isa CuArray
+  dv = gradient(v -> sum(abs2, re(v).z), cu([10, 20, 30, 40, 50.0]))[1]
+  @test dv isa CuArray
 end
 
 @testset "destructure mixed" begin