Don't use fix1 for enzyme (#4)

tpapp · Dec 24, 2022 · da73c64 · da73c64
1 parent 7c3bb0a
commit da73c64
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/AD_Enzyme.jl b/src/AD_Enzyme.jl
@@ -43,7 +43,8 @@ function logdensity_and_gradient(∇ℓ::EnzymeGradientLogDensity{<:Any,<:Enzyme
                                  x::AbstractVector)
     @unpack ℓ, mode, shadow = ∇ℓ
     _shadow = shadow === nothing ? Enzyme.onehot(x) : shadow
-    y, ∂ℓ_∂x = Enzyme.autodiff(mode, Base.Fix1(logdensity, ℓ), Enzyme.BatchDuplicated,
+    y, ∂ℓ_∂x = Enzyme.autodiff(mode, logdensity, Enzyme.BatchDuplicated,
+                               Enzyme.Const(ℓ),
                                Enzyme.BatchDuplicated(x, _shadow))
     return y, collect(∂ℓ_∂x)
 end
@@ -55,7 +56,8 @@ function logdensity_and_gradient(∇ℓ::EnzymeGradientLogDensity{<:Any,<:Enzyme
     # Ref: https://github.com/EnzymeAD/Enzyme.jl/issues/107
     y = logdensity(ℓ, x)
     ∂ℓ_∂x = zero(x)
-    Enzyme.autodiff(mode, Base.Fix1(logdensity, ℓ), Enzyme.Active,
+    Enzyme.autodiff(mode, logdensity, Enzyme.Active,
+                    Enzyme.Const(ℓ),
                     Enzyme.Duplicated(x, ∂ℓ_∂x))
     y, ∂ℓ_∂x
 end