Migrate docs build to GPU

Author: Avik Pal

.buildkite/pipeline.yml

@@ -0,0 +1,58 @@
+steps:
+  - label: ":julia: Julia {{matrix.julia}} + CUDA GPU"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.julia}}"
+      - JuliaCI/julia-test#v1:
+          test_args: "--quickfail"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+          dirs:
+            - src
+            - ext
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    env:
+      GROUP: "CUDA"
+    if: build.message !~ /\[skip tests\]/
+    timeout_in_minutes: 240
+    matrix:
+      setup:
+        julia:
+          - "1"
+
+  - label: ":open_book: Documentation"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+          dirs:
+            - src
+            - ext
+    command: |
+      julia --project --code-coverage=user --color=yes --threads=3 -e '
+        println("--- :julia: Instantiating project")
+        using Pkg
+        Pkg.instantiate()
+        Pkg.activate("docs")
+        Pkg.develop(PackageSpec(path=pwd()))
+        Pkg.instantiate()
+        println("+++ :julia: Building documentation")
+
+        Pkg.activate("docs")
+        include("docs/make.jl")'
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    env:
+      DATADEPS_ALWAYS_ACCEPT: true
+      JULIA_DEBUG: "Documenter"
+      GKSwstype: "100" # https://discourse.julialang.org/t/generation-of-documentation-fails-qt-qpa-xcb-could-not-connect-to-display/60988
+    if: build.message !~ /\[skip docs\]/
+    timeout_in_minutes: 240
+
+env:
+  SECRET_CODECOV_TOKEN: "fbSN+ZbScLIWr1FOpAu1Z8PYWFobqbLGFayOgZE1ebhE8LIH/PILGXUMcdm9gkXVSwgdETDD0s33k14lBkJ90O4dV9w6k79F/pEgzVHV8baMoXZG03BPMxztlcoRXrKtRtAp+MwoATc3Ldb9H5vqgAnVNn5rhn4Rp0Z6LOVRC43hbhKBBKYh/N4gqpIQlcW4dBXmELhlnMFnUILjwGRVgEt/zh8H+vmf0qiIulNIQ/rfGISROHqFML0QDL4icloiqX08J76ZP/gZCeg6rJ0gl3ok3IspNPz51rlbvijqsPNyIHWi29OrAtWX3qKHfrAOoGIrE1d5Oy4wx4XaN/YBhg==;U2FsdGVkX188gcRjkUNMEC2Z5fEFfhsYY4WJbhhINOuCUgqq9XNHVDbJhzFUFVQ+UiuPHFg7CW/gn+3IkSVyOA=="
+  SECRET_DOCUMENTER_KEY: "nIiMZCYpVq7ZJVZiSxS/xj1bqHWro1+5b2iq1cQluCDYEdWQc1QQu/PxfXUEaYYCj+1ZRQZabXv8Zqz2zdACeBT3btVgVf3EnWENJOK0BnP0XYpK/q1ARP1skexBIq3xokggvOcgf7n6JO3uuIwIQfbZK4oV7Cam92tJHmplUhbJdtAa63DlwzRgw4Ti5Brydl1QcIK4oX9PJmg1T/nybGFGD+IOtvaqUtxUJ9mlQhzfwITqFEyLI6ZpzAgzjQ84akiaz4ClHbpJSOZHR3ljSYkbrDiQA0Z2Cj9MZ84ApfNdA8SFIp66hnhkBEleX4g2bOFBX7FcZDT898roq1en9Q==;U2FsdGVkX18nuQu5WxxuFt5wZ2qgOXw7V2emSyef/iWOJYWHZA+vPceeUYoV8E2JCDbCdWoYJWylk9SUA8OvtibGv5D1lpUgKlGlc8YNWjDh3qDdNb04tFvmmmGvN/bAbMOtDaOuzH8wlduofdtC/OhV+TP6pmn2vMiaRM6xqrLhbZOjN77Cjb4KOLImwt2ejotu3WhKE2KJ2uwC24/rtfx9nQlh8HLSkJBZDM3+eQQTdJrmwJRXmXovm3Lvv4faDoRyc7sNHsS37MTUUwjAFSw8OunU27ia1/Y3CL7UC/oyvf8HnCWDBNL6wDUlImzUz2JZH0vyxvzT6X3IKL1roAPG9jZHk6AZ0rY5o3PfWkhMO2NuQE0AeCXUsIK6sNB8T1C+77EZjhu6hJSWqfaVdjKrmxL9xdCgIjQhcel1fh/svwTO5miB7RPhZxHBbSloSGt5siRSI31W/wOKhZKOIfA8EuRY1yNSf5nchbCS7gSCSXSfZ6eixsZr+gAgs7h+Pg5yARwMLV34sgvDwtKem6UZgOVzUc4WVp3IF5vP/aF4sI/bP5Yq1cTam42mjc3OJDSlQSebHygo+Aw24ovWsYccXQ+Q3Jt1j4rC+DuKRjTuU9NJYSZLRWh86QW0TAmdg4Ev5r6sOGD3vNqLZe9hxAOX6zIGaur8dIcruJHIDLDKdI9wVgFnM/hlFtuPnSvy961SvxAtIDxZUdS8Hhb083dwLK/j9UI/K47y6grkGq9451cmeRnCEfa3I+NQKkWOy6V0HAYymgA0MIS1Eca/1Q=="

.github/workflows/Documentation.yml

@@ -5,6 +5,7 @@
 
 [![codecov](https://codecov.io/gh/SciML/DeepEquilibriumNetworks.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/SciML/DeepEquilibriumNetworks.jl)
 [![Build Status](https://github.com/SciML/DeepEquilibriumNetworks.jl/workflows/CI/badge.svg)](https://github.com/SciML/DeepEquilibriumNetworks.jl/actions?query=workflow%3ACI)
+[![Build status](https://badge.buildkite.com/d7ce1858c4f89456c2d90e80c9b04b710bd81d7641db0a087c.svg?branch=main)](https://buildkite.com/julialang/deepequilibriumnetworks)
 
 [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor%27s%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
 [![SciML Code Style](https://img.shields.io/static/v1?label=code%20style&message=SciML&color=9558b2&labelColor=389826)](https://github.com/SciML/SciMLStyle)
@@ -24,8 +25,8 @@ Pkg.add("DeepEquilibriumNetworks")
 ## Quickstart
 
 ```julia
-using DeepEquilibriumNetworks, Lux, Random, NonlinearSolve, Zygote
-# using LuxCUDA, LuxAMDGPU ## Install and Load for GPU Support
+using DeepEquilibriumNetworks, Lux, Random, NonlinearSolve, Zygote, SciMLSensitivity
+# using LuxCUDA, LuxAMDGPU ## Install and Load for GPU Support. See https://lux.csail.mit.edu/dev/manual/gpu_management
 
 seed = 0
 rng = Random.default_rng()

README.md

@@ -3,6 +3,7 @@ DeepEquilibriumNetworks = "6748aba7-0e9b-415e-a410-ae3cc0ecb334"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

docs/Project.toml

@@ -16,9 +16,9 @@ Pkg.add("DeepEquilibriumNetworks")
 
 ## Quick-start
 
-```@example
+```@example quickstart
 using DeepEquilibriumNetworks, Lux, Random, NonlinearSolve, Zygote, SciMLSensitivity
-# using LuxCUDA, LuxAMDGPU ## Install and Load for GPU Support
+using LuxCUDA  # For NVIDIA GPU support
 
 seed = 0
 rng = Random.default_rng()
@@ -35,8 +35,11 @@ ps, st = Lux.setup(rng, model) |> gdev
 x = rand(rng, Float32, 2, 3) |> gdev
 y = rand(rng, Float32, 2, 3) |> gdev
 
-model(x, ps, st)
+y, 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))
 ```
 

docs/src/index.md

@@ -138,7 +138,7 @@ end
 
                 x = randn(rng, Float32, x_size...)
                 z, st = model(x, ps, st)
-                z_ = mapreduce(DEQs.__flatten, vcat, z)
+                z_ = DEQs.__flatten_vcat(z)
 
                 opt_broken = solver isa NewtonRaphson ||
                              solver isa SimpleLimitedMemoryBroyden ||
@@ -162,7 +162,7 @@ end
                 @test st.solution == DeepEquilibriumSolution()
 
                 z, st = model(x, ps, st)
-                z_ = mapreduce(DEQs.__flatten, vcat, z)
+                z_ = DEQs.__flatten_vcat(z)
                 opt_broken = jacobian_regularization isa AutoZygote
                 @jet model(x, ps, st) opt_broken=opt_broken