Initial version of CUDA working again

Author: Avik Pal

Manifest.toml

@@ -61,9 +61,9 @@ version = "0.2.0"
 
 [[deps.ArrayInterface]]
 deps = ["Adapt", "LinearAlgebra", "Requires", "SparseArrays", "SuiteSparse"]
-git-tree-sha1 = "247efbccf92448be332d154d6ca56b9fcdd93c31"
+git-tree-sha1 = "bbec08a37f8722786d87bedf84eae19c020c4efa"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "7.6.1"
+version = "7.7.0"
 
     [deps.ArrayInterface.extensions]
     ArrayInterfaceBandedMatricesExt = "BandedMatrices"
@@ -696,7 +696,7 @@ uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [[deps.LinearSolve]]
 deps = ["ArrayInterface", "ConcreteStructs", "DocStringExtensions", "EnumX", "FastLapackInterface", "GPUArraysCore", "InteractiveUtils", "KLU", "Krylov", "Libdl", "LinearAlgebra", "MKL_jll", "PrecompileTools", "Preferences", "RecursiveFactorization", "Reexport", "Requires", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Sparspak", "StaticArraysCore", "UnPack"]
-git-tree-sha1 = "7c9f62dc6f7d11b0f8dfb3c1f3637cb220dc9866"
+git-tree-sha1 = "b1148a6596c5cd9b2d848c26b500c79d102ffc5d"
 repo-rev = "ap/normal_cholesky_dispatches"
 repo-url = "https://github.com/SciML/LinearSolve.jl.git"
 uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
@@ -1248,7 +1248,7 @@ version = "0.3.7"
 
 [[deps.SciMLSensitivity]]
 deps = ["ADTypes", "Adapt", "ArrayInterface", "ChainRulesCore", "DiffEqBase", "DiffEqCallbacks", "DiffEqNoiseProcess", "Distributions", "EllipsisNotation", "Enzyme", "FiniteDiff", "ForwardDiff", "FunctionProperties", "FunctionWrappersWrappers", "Functors", "GPUArraysCore", "LinearAlgebra", "LinearSolve", "Markdown", "OrdinaryDiffEq", "Parameters", "PreallocationTools", "QuadGK", "Random", "RandomNumbers", "RecursiveArrayTools", "Reexport", "ReverseDiff", "SciMLBase", "SciMLOperators", "SparseDiffTools", "StaticArrays", "StaticArraysCore", "Statistics", "StochasticDiffEq", "Tracker", "TruncatedStacktraces", "Zygote"]
-git-tree-sha1 = "71efe41945d384c3b4873ba681ef8905cae79123"
+git-tree-sha1 = "4120f3ef35508d2f9d967f40364765da4950cefe"
 repo-rev = "ap/ssadjoint_fix"
 repo-url = "https://github.com/SciML/SciMLSensitivity.jl.git"
 uuid = "1ed8b502-d754-442c-8d5d-10ac956f44a1"

docs/src/index.md

@@ -35,12 +35,12 @@ ps, st = Lux.setup(rng, model) |> gdev
 x = rand(rng, Float32, 2, 3) |> gdev
 y = rand(rng, Float32, 2, 3) |> gdev
 
-y, st_ = model(x, ps, st)
+res, st_ = model(x, ps, st)
 st_.layer_2.solution
 ```
 
 ```@example quickstart
-gs = only(Zygote.gradient(p -> sum(abs2, first(first(model(x, p, st))) .- y), ps))
+gs = only(Zygote.gradient(p -> sum(abs2, first(model(x, p, st)) .- y), ps))
 ```
 
 ## Citation

src/layers.jl

@@ -308,7 +308,7 @@ julia> model(x, ps, st);
 function MultiScaleDeepEquilibriumNetwork(main_layers::Tuple, mapping_layers::Matrix,
         post_fuse_layer::Union{Nothing, Tuple}, solver, scales;
         jacobian_regularization=nothing, kwargs...)
-    @assert jacobian_regularization === nothing "Jacobian Regularization is not supported yet for MultiScale Models."
+    @assert jacobian_regularization===nothing "Jacobian Regularization is not supported yet for MultiScale Models."
     l1 = Parallel(nothing, main_layers...)
     l2 = BranchLayer(Parallel.(+, map(x -> tuple(x...), eachrow(mapping_layers))...)...)
 

test/layers.jl

@@ -11,7 +11,7 @@ function loss_function(model, x, ps, st)
     return l1 + l2 + l3
 end
 
-@testset "DeepEquilibriumNetwork" begin
+@testset "DeepEquilibriumNetwork: $(mode)" for (mode, aType, dev, ongpu) in MODES
     rng = __get_prng(0)
 
     base_models = [
@@ -41,10 +41,10 @@ end
                 SkipDeepEquilibriumNetwork(base_model, solver; jacobian_regularization)
             end
 
-            ps, st = Lux.setup(rng, model)
+            ps, st = Lux.setup(rng, model) |> dev
             @test st.solution == DeepEquilibriumSolution()
 
-            x = randn(rng, Float32, x_size...)
+            x = randn(rng, Float32, x_size...) |> dev
             z, st = model(x, ps, st)
 
             opt_broken = solver isa NewtonRaphson ||
@@ -62,7 +62,7 @@ end
             @test __is_finite_gradient(gs_x)
             @test __is_finite_gradient(gs_ps)
 
-            ps, st = Lux.setup(rng, model)
+            ps, st = Lux.setup(rng, model) |> dev
             st = Lux.update_state(st, :fixed_depth, Val(10))
             @test st.solution == DeepEquilibriumSolution()
 
@@ -82,7 +82,7 @@ end
     end
 end
 
-@testset "MultiScaleDeepEquilibriumNetworks" begin
+@testset "MultiScaleDeepEquilibriumNetwork: $(mode)" for (mode, aType, dev, ongpu) in MODES
     rng = __get_prng(0)
 
     main_layers = [

test/runtests.jl

@@ -1,8 +1,5 @@
 using SafeTestsets, Test
 
-# TODO: CUDA Testing
-const GROUP = get(ENV, "GROUP", "ALL")
-
 @testset "Deep Equilibrium Networks" begin
     @safetestset "Quality Assurance" begin
         include("qa.jl")

test/test_utils.jl

@@ -1,5 +1,6 @@
 using DeepEquilibriumNetworks, Functors, Lux, Random, StableRNGs, Zygote
 import LuxTestUtils: @jet
+using LuxCUDA
 
 __nameof(::X) where {X} = nameof(X)
 
@@ -26,3 +27,21 @@ function __get_conv_layer(args...; kwargs...)
     init_weight(rng::AbstractRNG, dims...) = randn(rng, Float32, dims) .* 0.001f0
     return Conv(args...; init_weight, use_bias=false, kwargs...)
 end
+
+const GROUP = get(ENV, "GROUP", "All")
+
+cpu_testing() = GROUP == "All" || GROUP == "CPU"
+cuda_testing() = LuxCUDA.functional() && (GROUP == "All" || GROUP == "CUDA")
+
+if !@isdefined(MODES)
+    const MODES = begin
+        cpu_mode = ("CPU", Array, LuxCPUDevice(), false)
+        cuda_mode = ("CUDA", CuArray, LuxCUDADevice(), true)
+
+        modes = []
+        cpu_testing() && push!(modes, cpu_mode)
+        cuda_testing() && push!(modes, cuda_mode)
+
+        modes
+    end
+end

test/utils.jl

@@ -2,10 +2,10 @@ using DeepEquilibriumNetworks, LinearAlgebra, SciMLBase, Test
 
 include("test_utils.jl")
 
-@testset "split_and_reshape" begin
-    x1 = ones(Float32, 4, 4)
-    x2 = fill(0.5f0, 2, 4)
-    x3 = zeros(Float32, 1, 4)
+@testset "split_and_reshape: $mode" for (mode, aType, dev, ongpu) in MODES
+    x1 = ones(Float32, 4, 4) |> aType
+    x2 = fill(0.5f0, 2, 4) |> aType
+    x3 = zeros(Float32, 1, 4) |> aType
 
     x = vcat(x1, x2, x3)
     split_idxs = Val(cumsum((0, size(x1, 1), size(x2, 1), size(x3, 1))))