feat: support runtime activity for enzyme

CHANGELOG.md

@@ -5,6 +5,13 @@ 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.3.0] - 2024-09-22
+
+### Added
+
+  - Adds a kwarg `enzyme_set_runtime_activity` to `test_gradients` to allow users to set
+    the runtime activity of Enzyme tests.
+
 ## [1.2.0] - 2024-09-18
 
 ### Added

Project.toml

@@ -1,7 +1,7 @@
 name = "LuxTestUtils"
 uuid = "ac9de150-d08f-4546-94fb-7472b5760531"
 authors = ["Avik Pal <[email protected]>"]
-version = "1.2.1"
+version = "1.3.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/autodiff.jl

@@ -114,6 +114,7 @@ Test the gradients of `f` with respect to `args` using the specified backends.
   - `soft_fail`: If `true`, then the test will be recorded as a `soft_fail` test. This
     overrides any `broken` kwargs. Alternatively, a list of backends can be passed to
     `soft_fail` to allow soft_fail tests for only those backends.
+  - `enzyme_set_runtime_activity`: If `true`, then activate runtime activity for Enzyme.
   - `kwargs`: Additional keyword arguments to pass to `check_approx`.
 
 ## Example
@@ -129,6 +130,7 @@ julia> test_gradients(f, 1.0, x, nothing)
 """
 function test_gradients(f, args...; skip_backends=[], broken_backends=[],
         soft_fail::Union{Bool, Vector}=false,
+        enzyme_set_runtime_activity::Bool=false,
         # Internal kwargs start
         source::LineNumberNode=LineNumberNode(0, nothing),
         test_expr::Expr=:(check_approx(∂args, ∂args_gt; kwargs...)),
@@ -146,7 +148,12 @@ function test_gradients(f, args...; skip_backends=[], broken_backends=[],
         total_length ≤ 100 && push!(backends, AutoForwardDiff())
         total_length ≤ 100 && push!(backends, AutoFiniteDiff())
         # TODO: Move Enzyme out of here once it supports GPUs
-        ENZYME_TESTING_ENABLED && push!(backends, AutoEnzyme())
+        if ENZYME_TESTING_ENABLED
+            mode = enzyme_set_runtime_activity ?
+                   Enzyme.set_runtime_activity(Enzyme.Reverse) :
+                   Enzyme.Reverse
+            push!(backends, AutoEnzyme(; mode))
+        end
     end
     push!(backends, AutoTracker())