feat: better test integration in test_gradients

LuxDL · Sep 18, 2024 · 3f3cd9f · 3f3cd9f
1 parent 6e42fd2
commit 3f3cd9f
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 19 deletions.
diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml
@@ -5,4 +5,3 @@ indent = 4
 format_docstrings = true
 separate_kwargs_with_semicolon = true
 always_for_in = true
-join_lines_based_on_source = false
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -27,6 +27,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
+          - "min"
           - "1"
           - "pre"
         os:
@@ -64,7 +65,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     timeout-minutes: 60
     env:
-      GROUP: ${{ matrix.package.group }}
+      BACKEND_GROUP: ${{ matrix.package.group }}
     strategy:
       fail-fast: false
       matrix:
@@ -126,8 +127,6 @@ jobs:
       - uses: julia-actions/julia-downgrade-compat@v1
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
-        env:
-          LUX_TEST_GROUP: ${{ matrix.test_group }}
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v4
         with:

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,12 @@ All notable changes to this project since the release of v1 will be documented i
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.2.0] - 2024-09-17
+
+### Added
+
+  - By default, we no longer wrap the entire gradient computation in a `@test` macro.
+
 ## [1.1.4] - 2024-08-21
 
 ### Fixed

diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxTestUtils"
 uuid = "ac9de150-d08f-4546-94fb-7472b5760531"
 authors = ["Avik Pal <avikpal@mit.edu>"]
-version = "1.1.4"
+version = "1.2.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/src/autodiff.jl b/src/autodiff.jl
@@ -128,7 +128,13 @@ julia> test_gradients(f, 1.0, x, nothing)
 ```
 """
 function test_gradients(f, args...; skip_backends=[], broken_backends=[],
-        soft_fail::Union{Bool, Vector}=false, kwargs...)
+        soft_fail::Union{Bool, Vector}=false,
+        # Internal kwargs start
+        source=LineNumberNode(0, nothing),
+        test_expr=:(check_approx(∂args, ∂args_gt; kwargs...)),
+        # Internal kwargs end
+        kwargs...)
+    # TODO: We should add a macro version that propagates the line number info and the test_expr
     on_gpu = get_device_type(args) <: AbstractGPUDevice
     total_length = mapreduce(__length, +, Functors.fleaves(args); init=0)
 
@@ -157,36 +163,58 @@ function test_gradients(f, args...; skip_backends=[], broken_backends=[],
 
     @testset "gradtest($(f))" begin
         @testset "$(nameof(typeof(backends[1])))() vs $(nameof(typeof(backend)))()" for backend in backends[2:end]
-            if backend in skip_backends
-                @test_skip begin
-                    ∂args = allow_unstable() do
-                        return gradient(f, backend, args...)
-                    end
-                    check_approx(∂args, ∂args_gt; kwargs...)
-                end
+            local_test_expr = :([$(nameof(typeof(backend)))] - $(test_expr))
+
+            result = if backend in skip_backends
+                Broken(:skipped, local_test_expr)
             elseif (soft_fail isa Bool && soft_fail) ||
                    (soft_fail isa Vector && backend in soft_fail)
-                @test_softfail begin
+                try
                     ∂args = allow_unstable() do
                         return gradient(f, backend, args...)
                     end
-                    check_approx(∂args, ∂args_gt; kwargs...)
+                    matched = check_approx(∂args, ∂args_gt; kwargs...)
+                    if matched
+                        Pass(:test, local_test_expr, nothing, nothing, source)
+                    else
+                        Broken(:test, local_test_expr)
+                    end
+                catch
+                    Broken(:test, local_test_expr)
                 end
             elseif backend in broken_backends
-                @test_broken begin
+                try
                     ∂args = allow_unstable() do
                         return gradient(f, backend, args...)
                     end
-                    check_approx(∂args, ∂args_gt; kwargs...)
+                    matched = check_approx(∂args, ∂args_gt; kwargs...)
+                    if matched
+                        Error(:test_unbroken, local_test_expr, matched, nothing, source)
+                    else
+                        Broken(:test, local_test_expr)
+                    end
+                catch
+                    Broken(:test, local_test_expr)
                 end
             else
-                @test begin
+                try
                     ∂args = allow_unstable() do
                         return gradient(f, backend, args...)
                     end
-                    check_approx(∂args, ∂args_gt; kwargs...)
+                    matched = check_approx(∂args, ∂args_gt; kwargs...)
+                    if matched
+                        Pass(:test, local_test_expr, nothing, nothing, source)
+                    else
+                        context = "\n   ∂args: $(∂args)\n∂args_gt: $(∂args_gt)"
+                        Fail(
+                            :test, local_test_expr, matched, nothing, context, source, false)
+                    end
+                catch err
+                    err isa InterruptException && rethrow()
+                    Error(:test, local_test_expr, err, Base.current_exceptions(), source)
                 end
             end
+            Test.record(get_testset(), result)
         end
     end
 end