Skip to content

Add defer_within_autodiff to EnzymeInterpreter#2254

Merged
wsmoses merged 2 commits intoEnzymeAD:mainfrom
jumerckx:deferred_within_autodiff
Aug 31, 2025
Merged

Add defer_within_autodiff to EnzymeInterpreter#2254
wsmoses merged 2 commits intoEnzymeAD:mainfrom
jumerckx:deferred_within_autodiff

Conversation

@jumerckx
Copy link
Copy Markdown
Contributor

@jumerckx jumerckx commented Jan 7, 2025

Together with Reactant pr: EnzymeAD/Reactant.jl#490

Comment thread src/compiler/interpreter.jl Outdated
Comment thread src/compiler/interpreter.jl Outdated
(; fargs, argtypes) = arginfo

if f === Enzyme.within_autodiff
if !(interp.defer_within_autodiff) && f === Enzyme.within_autodiff
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this necessary? This fundamentally breaks this functionality?

Copy link
Copy Markdown
Contributor Author

@jumerckx jumerckx Jan 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is for a Reactant bug: EnzymeAD/Reactant.jl#442 (comment)
Reason being that Reactant uses EnzymeInterpreter as well, while not necessarily doing autodiff.

@codecov-commenter
Copy link
Copy Markdown

codecov-commenter commented Jan 7, 2025

Codecov Report

❌ Patch coverage is 33.33333% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.92%. Comparing base (4ba9b71) to head (f1e15c9).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
src/compiler/interpreter.jl 33.33% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2254      +/-   ##
==========================================
- Coverage   74.93%   74.92%   -0.01%     
==========================================
  Files          56       56              
  Lines       17434    17436       +2     
==========================================
  Hits        13064    13064              
- Misses       4370     4372       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@wsmoses
Copy link
Copy Markdown
Member

wsmoses commented Jan 16, 2025

@vchuravy can you give this a review before merge

@vchuravy
Copy link
Copy Markdown
Member

Seems fine.

@vchuravy vchuravy force-pushed the deferred_within_autodiff branch from 3e37885 to 9f54068 Compare January 17, 2025 19:44
@avik-pal
Copy link
Copy Markdown
Collaborator

avik-pal commented Mar 7, 2025

bump on this

…_autodiff` to no return true during Reactant compilation.

When this flag is true, `interp.handler` is responsible for handling within_autodiff, or to toggle defer_within_autodiff to false somewhere down the call chain.
@wsmoses wsmoses force-pushed the deferred_within_autodiff branch from 9f54068 to f1e15c9 Compare August 31, 2025 18:19
@github-actions
Copy link
Copy Markdown
Contributor

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.
diff --git a/src/compiler/interpreter.jl b/src/compiler/interpreter.jl
index 77f0027..260c90b 100644
--- a/src/compiler/interpreter.jl
+++ b/src/compiler/interpreter.jl
@@ -173,7 +173,7 @@ function EnzymeInterpreter(
     reverse_rules::Bool,
     inactive_rules::Bool,
     broadcast_rewrite::Bool = true,
-    within_autodiff_rewrite::Bool = true,
+        within_autodiff_rewrite::Bool = true,
     handler = nothing
 )
     @assert world <= Base.get_world_counter()
@@ -250,23 +250,27 @@ EnzymeInterpreter(
     handler = nothing
 ) = EnzymeInterpreter(cache_or_token, mt, world, mode == API.DEM_ForwardMode, mode == API.DEM_ReverseModeCombined || mode == API.DEM_ReverseModePrimal || mode == API.DEM_ReverseModeGradient, inactive_rules, broadcast_rewrite, within_autodiff_rewrite, handler)
 
-function EnzymeInterpreter(interp::EnzymeInterpreter;
-    cache_or_token = (@static if HAS_INTEGRATED_CACHE
-        interp.token
-    else
-        interp.code_cache
-    end),
-    mt = interp.method_table,
-    local_cache = interp.local_cache,
-    world = interp.world,
-    inf_params = interp.inf_params,
-    opt_params = interp.opt_params,
-    forward_rules = interp.forward_rules,
-    reverse_rules = interp.reverse_rules,
-    inactive_rules = interp.inactive_rules,
-    broadcast_rewrite = interp.broadcast_rewrite,
-    within_autodiff_rewrite = interp.within_autodiff_rewrite,
-    handler = interp.handler)
+function EnzymeInterpreter(
+        interp::EnzymeInterpreter;
+        cache_or_token = (
+            @static if HAS_INTEGRATED_CACHE
+                interp.token
+            else
+                interp.code_cache
+            end
+        ),
+        mt = interp.method_table,
+        local_cache = interp.local_cache,
+        world = interp.world,
+        inf_params = interp.inf_params,
+        opt_params = interp.opt_params,
+        forward_rules = interp.forward_rules,
+        reverse_rules = interp.reverse_rules,
+        inactive_rules = interp.inactive_rules,
+        broadcast_rewrite = interp.broadcast_rewrite,
+        within_autodiff_rewrite = interp.within_autodiff_rewrite,
+        handler = interp.handler
+    )
     return EnzymeInterpreter(
         cache_or_token,
         mt,

@wsmoses wsmoses merged commit 1aa2002 into EnzymeAD:main Aug 31, 2025
24 of 31 checks passed
@jumerckx jumerckx deleted the deferred_within_autodiff branch November 17, 2025 19:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants