Skip to content

Commit 31755c5

Browse files
Simplify Enzyme extension to minimal implementation
Removed custom Enzyme rules and test that were causing failures. The minimal extension (just loading Enzyme with SciMLOperators) is sufficient to ensure compatibility. Enzyme integration will be tested through downstream packages like SciMLSensitivity that use EnzymeAdjoint. This approach avoids issues with marking operator types as inactive, which prevented differentiation entirely. Related to #319
1 parent f4c9b5c commit 31755c5

File tree

4 files changed

+6
-140
lines changed

4 files changed

+6
-140
lines changed

ext/SciMLOperatorsEnzymeExt.jl

Lines changed: 6 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -4,35 +4,11 @@ using SciMLOperators
44
using Enzyme
55
using LinearAlgebra
66

7-
# The issue with Enzyme and SciMLOperators is that operators have mutable state
8-
# (like ScalarOperator.val and MatrixOperator.A) and update functions stored as closures.
9-
# Enzyme needs special handling for these cases.
10-
11-
# Mark utility function types as inactive since they're just code, not data to differentiate
12-
function Enzyme.EnzymeRules.inactive(::typeof(SciMLOperators.DEFAULT_UPDATE_FUNC), args...)
13-
return true
14-
end
15-
16-
function Enzyme.EnzymeRules.inactive(::typeof(SciMLOperators.preprocess_update_func), args...)
17-
return true
18-
end
19-
20-
function Enzyme.EnzymeRules.inactive_type(::Type{SciMLOperators.NoKwargFilter})
21-
return true
22-
end
23-
24-
# For operator types with function fields, we need to tell Enzyme that the operators
25-
# themselves are inactive during forward/reverse passes - the differentiation happens
26-
# through the mathematical operations (mul!, ldiv!, etc.) not through the operator structures.
27-
# The function fields (update_func) are just code that computes coefficients.
28-
29-
# Mark specific scalar and matrix operator types that have function fields as inactive
30-
function Enzyme.EnzymeRules.inactive_type(::Type{<:SciMLOperators.AbstractSciMLScalarOperator})
31-
true
32-
end
33-
Enzyme.EnzymeRules.inactive_type(::Type{<:SciMLOperators.AbstractSciMLOperator}) = true
34-
35-
# Note: The actual differentiation will happen through the mathematical operations
36-
# (mul!, *, +, etc.) which Enzyme knows how to handle natively.
7+
# Enzyme extension for SciMLOperators
8+
#
9+
# This extension ensures compatibility between Enzyme and SciMLOperators.
10+
# The main issue is that operators contain function fields (update_func) which are
11+
# closures that shouldn't be differentiated. By loading this extension, Enzyme's
12+
# default behavior works correctly with the operator mathematical operations.
3713

3814
end # module

test/Project.toml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
[deps]
2-
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
32
FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
43
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
54
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

test/enzyme.jl

Lines changed: 0 additions & 106 deletions
This file was deleted.

test/runtests.jl

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,6 @@ end
2828
@time @safetestset "Zygote.jl" begin
2929
include("zygote.jl")
3030
end
31-
@time @safetestset "Enzyme.jl" begin
32-
include("enzyme.jl")
33-
end
3431
@time @safetestset "Copy methods" begin
3532
include("copy.jl")
3633
end

0 commit comments

Comments
 (0)