Start adding tests

.JuliaFormatter.toml

@@ -1,6 +1,5 @@
 style = "sciml"
 whitespace_in_kwargs = false
-always_use_return = true
 margin = 92
 indent = 4
 format_docstrings = true

.buildkite/pipeline.yml

@@ -1,50 +1,32 @@
 steps:
-  - label: ":julia: Format Check"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "1"
-    agents:
-      queue: "default"
-    if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 240
-    command: |
-      julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))'
-      julia  -e 'using JuliaFormatter; format(".", verbose=true)'
-      julia -e '
-          out = Cmd(`git diff --name-only`) |> read |> String
-          if out == ""
-              exit(0)
-          else
-              @error "Some files have not been formatted !!!"
-              write(stdout, out)
-              exit(1)
-          end'
-
-  - label: ":julia: Julia {{matrix.julia}} (CPU)"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "{{matrix.julia}}"
-      - JuliaCI/julia-test#v1:
-          test_args: "--quickfail"
-      - JuliaCI/julia-coverage#v1:
-          codecov: true
-          dirs:
-            - src
-    agents:
-      queue: "default"
-    env:
-      GROUP: "CPU"
-    if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 240
-    matrix:
-      setup:
-        julia:
-          - "1"
-          - "nightly"
-      adjustments:
-        - with:
-            julia: "nightly"
-          soft_fail: true
+  - group: ":julia: CUDA GPU"
+    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"
 
 env:
-  CODECOV_TOKEN: "9349e5e6-16b3-4c9c-9496-304a84a7358d" # Someday I will regret doing this
+  RETESTITEMS_NWORKERS: 4
+  RETESTITEMS_NWORKER_THREADS: 2
+  JULIA_AMDGPU_LOGGING_ENABLED: true
+  RETESTITEMS_TESTITEM_TIMEOUT: 10000
+  SECRET_CODECOV_TOKEN: "NkRNr3aGg3k4bKi07RGhGdhhsPV8t97y0VASfmra5BzT+6h/S+hEZ6p7U6SE0/1LQrHxRBy9vaWiwF+VW1ZHk7KMUetuOYmymXON/AUBbiE4LsfFOVwrna7U0kuqWHZbdKn8XAJxu6au1uRMrOXPXw176KkuWRzwF/jLWvvv7s+KqX4oaiDirXxGCRSssVizT2hdWkkrtct+GjLeF/g9jgGa8xn8j2Pp8AS62EPMoC/YKgV/e3yK58LSOKOBF+1ddvYzaFoDABkNMehHA52MNXgDoxikTc0YGnd8nMGfTUiRPaNLHRQaXS/M0oaVT7PkXFlJe6O6izCnkIx2+Ix57w==;U2FsdGVkX1/k+16T0rj/Tntc6gOaH8GwRDvncs1a+BwbnnnnXeAmiwvbowfRSnoldKtpHhJcwLQLbXFDXD8U6g=="

.github/workflows/CI.yml

@@ -37,10 +37,15 @@ jobs:
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         env:
-            GROUP: "CPU"
+          GROUP: "CPU"
+          RETESTITEMS_NWORKERS: 4
+          RETESTITEMS_NWORKER_THREADS: 2
       - uses: julia-actions/julia-processcoverage@v1
         with:
-          directories: src
-      - uses: codecov/codecov-action@v3
+          directories: src,ext
+      - uses: codecov/codecov-action@v4
         with:
           files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          verbose: true
+          fail_ci_if_error: true

.github/workflows/CompatHelper.yml

@@ -0,0 +1,41 @@
+name: CompatHelper
+on:
+  schedule:
+    - cron: 0 0 * * *
+  workflow_dispatch:
+permissions:
+  contents: write
+  pull-requests: write
+jobs:
+  CompatHelper:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v1
+        with:
+          version: 1
+      - name: "Add the General registry via Git"
+        run: |
+          import Pkg
+          ENV["JULIA_PKG_SERVER"] = ""
+          Pkg.Registry.add("General")
+        shell: julia --color=yes {0}
+      - name: "Install CompatHelper"
+        run: |
+          import Pkg
+          name = "CompatHelper"
+          uuid = "aa819f21-2bde-4658-8897-bab36330d9b7"
+          version = "3"
+          Pkg.add(; name, uuid, version)
+        shell: julia --color=yes {0}
+      - name: "Run CompatHelper"
+        run: |
+          import CompatHelper
+          subdirs = [""]
+          append!(subdirs, joinpath.(("examples",), filter(p -> isdir(joinpath("examples", p)), readdir("examples"))))
+          CompatHelper.main(; subdirs)
+        shell: julia --color=yes {0}
+        working-directory: "./"
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        #   COMPATHELPER_PRIV: ${{ secrets.DOCUMENTER_KEY }}

.github/workflows/Downgrade.yml

@@ -0,0 +1,41 @@
+name: Downgrade
+on:
+  pull_request:
+    branches:
+      - main
+    paths-ignore:
+      - 'docs/**'
+  push:
+    branches:
+      - main
+    paths-ignore:
+      - 'docs/**'
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        version: ['1.9']
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v1
+        with:
+          version: ${{ matrix.version }}
+      - uses: cjdoris/julia-downgrade-compat-action@v1
+        with:
+          skip: Pkg,TOML
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+        env:
+          GROUP: "CPU"
+          RETESTITEMS_NWORKERS: 4
+          RETESTITEMS_NWORKER_THREADS: 2
+      - uses: julia-actions/julia-processcoverage@v1
+        with:
+          directories: src,ext
+      - uses: codecov/codecov-action@v4
+        with:
+          files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          verbose: true
+          fail_ci_if_error: true

.github/workflows/FormatPR.yml

@@ -26,4 +26,4 @@ jobs:
       - name: Check outputs
         run: |
           echo "Pull Request Number - ${{ steps.cpr.outputs.pull-request-number }}"
-          echo "Pull Request URL - ${{ steps.cpr.outputs.pull-request-url }}"
+          echo "Pull Request URL - ${{ steps.cpr.outputs.pull-request-url }}"

.github/workflows/TagBot.yml

@@ -0,0 +1,31 @@
+name: TagBot
+on:
+  issue_comment:
+    types:
+      - created
+  workflow_dispatch:
+    inputs:
+      lookback:
+        default: "3"
+permissions:
+  actions: read
+  checks: read
+  contents: write
+  deployments: read
+  issues: read
+  discussions: read
+  packages: read
+  pages: read
+  pull-requests: read
+  repository-projects: read
+  security-events: read
+  statuses: read
+jobs:
+  TagBot:
+    if: github.event_name == 'workflow_dispatch' || github.actor == 'JuliaTagBot'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: JuliaRegistries/TagBot@v1
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}
+        #   ssh: ${{ secrets.DOCUMENTER_KEY }}

LocalPreferences.toml

@@ -0,0 +1,2 @@
+[LuxTestUtils]
+target_modules = ["BatchedRoutines"]

Project.toml

@@ -31,22 +31,41 @@ BatchedRoutinesZygoteExt = ["Zygote"]
 
 [compat]
 ADTypes = "0.2.6"
+Adapt = "4.0.3"
+Aqua = "0.8.4"
 ArrayInterface = "7.8.1"
-ChainRulesCore = "1.23.0"
+ChainRulesCore = "1.23"
+ConcreteStructs = "0.2.3"
+ExplicitImports = "1.4.0"
 FastClosures = "0.3.2"
-FiniteDiff = "2.22.0"
+FillArrays = "1.9.3"
+FiniteDiff = "2.22"
 ForwardDiff = "0.10.36"
 LinearAlgebra = "1.10"
+LuxDeviceUtils = "0.1.17"
+LuxTestUtils = "0.1.15"
 PrecompileTools = "1.2.0"
-ReverseDiff = "1.15.1"
+Random = "<0.0.1, 1"
+ReTestItems = "1.23.1"
+ReverseDiff = "1.15"
+StableRNGs = "1.0.1"
+Zygote = "0.6.69"
 julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+LuxDeviceUtils = "34f89e08-e1d5-43b4-8944-0b49ac560553"
+LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "ExplicitImports", "FiniteDiff", "ForwardDiff", "ReTestItems"]
+test = ["Aqua", "ExplicitImports", "FiniteDiff", "ForwardDiff", "LuxCUDA", "LuxDeviceUtils", "LuxTestUtils", "Random", "ReTestItems", "ReverseDiff", "StableRNGs", "Test", "Zygote"]

src/chainrules.jl

@@ -15,7 +15,7 @@ function CRC.rrule(::typeof(batched_jacobian), ad, f::F, x::AbstractMatrix) wher
     ∇batched_jacobian = Δ -> begin
         gradient_ad = AutoZygote()
         _map_fnₓ = ((i, Δᵢ),) -> _jacobian_vector_product(AutoForwardDiff(),
-            x -> batched_gradient(gradient_ad, x_ -> sum(vec(f(x_, p))[i:i]), x),
+            x -> batched_gradient(gradient_ad, x_ -> sum(vec(f(x_))[i:i]), x),
             x, reshape(Δᵢ, size(x)))
         ∂x = reshape(mapreduce(_map_fnₓ, +, enumerate(eachrow(Δ))), size(x))
         return NoTangent(), NoTangent(), NoTangent(), ∂x

src/matrix.jl

@@ -125,7 +125,7 @@ function Base.Matrix(A::UniformBlockDiagonalMatrix)
     L1, L2, _ = size(A.data)
     fill!(M, false)
     for (i, Aᵢ) in enumerate(batchview(A))
-        M[((i - 1) * L1 + 1):(i * L1), ((i - 1) * L2 + 1):(i * L2)] .= Aᵢ
+        M[((i - 1) * L1 + 1):(i * L1), ((i - 1) * L2 + 1):(i * L2)] .= Matrix(Aᵢ)
     end
     return M
 end
@@ -307,7 +307,9 @@ function Base.size(F::GenericBatchedFactorization, i::Integer)
 end
 Base.eltype(F::GenericBatchedFactorization) = eltype(first(F.fact))
 
-LinearAlgebra.issuccess(fact::GenericBatchedFactorization) = all(issuccess, fact.fact)
+function LinearAlgebra.issuccess(fact::GenericBatchedFactorization)
+    return all(LinearAlgebra.issuccess, fact.fact)
+end
 
 function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, F::GenericBatchedFactorization)
     println(io, "GenericBatchedFactorization() with Batch Count: $(nbatches(F))")

test/autodiff_tests.jl

@@ -1,3 +1,39 @@
-@testitem "AutoDiff Testing" begin
-    function simple_batched_funct() end
+@testitem "Batched Jacobians" setup=[SharedTestSetup] begin
+    using FiniteDiff, ForwardDiff, ReverseDiff, Zygote
+
+    rng = get_stable_rng(1001)
+
+    @testset "$mode" for (mode, aType, device, ongpu) in MODES
+        simple_batched_function = function (X, p)
+            X_ = reshape(X, :, nbatches(X))
+            return sum(abs2, X_ .* p; dims=1) .- sum(abs, X_ .* p; dims=1) .+ p .^ 2
+        end
+
+        X = randn(rng, 3, 2, 4) |> aType
+        p = randn(rng, 6) |> aType
+
+        J_fdiff = batched_jacobian(
+            AutoFiniteDiff(), simple_batched_function, Array(X), Array(p))
+        J_fwdiff = batched_jacobian(AutoForwardDiff(), simple_batched_function, X, p)
+
+        @test Matrix(J_fdiff)≈Matrix(J_fwdiff) atol=1e-3
+
+        X = randn(rng, 2, 4) |> aType
+        p = randn(rng, 2) |> aType
+
+        J_fdiff = batched_jacobian(
+            AutoFiniteDiff(), simple_batched_function, Array(X), Array(p))
+        J_fwdiff = batched_jacobian(AutoForwardDiff(), simple_batched_function, X, p)
+
+        @test Matrix(J_fdiff)≈Matrix(J_fwdiff) atol=1e-3
+
+        X = randn(rng, 3) |> aType
+        p = randn(rng, 3) |> aType
+
+        J_fdiff = batched_jacobian(
+            AutoFiniteDiff(), simple_batched_function, Array(X), Array(p))
+        J_fwdiff = batched_jacobian(AutoForwardDiff(), simple_batched_function, X, p)
+
+        @test Matrix(J_fdiff)≈Matrix(J_fwdiff) atol=1e-3
+    end
 end

test/qa_tests.jl

@@ -0,0 +1,7 @@
+@testitem "Quality Assurance" setup=[SharedTestSetup] begin
+    using Aqua, ExplicitImports
+
+    Aqua.test_all(BatchedRoutines; ambiguities=false)
+    Aqua.test_ambiguities(BatchedRoutines; broken=true)
+    @test ExplicitImports.check_no_implicit_imports(BatchedRoutines) === nothing
+end

test/runtests.jl

@@ -1,3 +1,3 @@
-using ReTestItems: ReTestItems
+using ReTestItems
 
 ReTestItems.runtests(@__DIR__)