fix CI and clean up

ACEsuit · Sep 29, 2023 · 7f53212 · 7f53212
1 parent 6c56778
commit 7f53212
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 600 deletions.
diff --git a/src/backflowpooling.jl b/src/backflowpooling.jl
@@ -1,11 +1,13 @@
 using ACEpsi.AtomicOrbitals: AtomicOrbitalsBasisLayer
 using LuxCore: AbstractExplicitLayer
 using Random: AbstractRNG
-using ChainRulesCore
 using ChainRulesCore: NoTangent
+
 using Polynomials4ML: _make_reqfields, @reqfields, POOL, TMP, META
 using ObjectPools: acquire!
 
+import ChainRulesCore: rrule
+
 mutable struct BackflowPooling
    basis::AtomicOrbitalsBasisLayer
    @reqfields
@@ -86,7 +88,7 @@ function evaluate(pooling::BackflowPooling, ϕnlm::AbstractArray, Σ::AbstractVe
 end
 
 # --------------------- connect with ChainRule
-function ChainRulesCore.rrule(::typeof(evaluate), pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) 
+function rrule(::typeof(evaluate), pooling::BackflowPooling, ϕnlm, Σ::AbstractVector) 
    A = pooling(ϕnlm, Σ)
    function pb(∂A)
       return NoTangent(), NoTangent(), _pullback_evaluate(∂A, pooling, ϕnlm, Σ), NoTangent()