rebase

FluxML · Oct 8, 2023 · 1933a9a · 1933a9a
1 parent 597bcd7
commit 1933a9a
Show file tree

Hide file tree

Showing 4 changed files with 2 additions and 384 deletions.
diff --git a/ext/NNlibCUDACUDNNExt/batchnorm.jl b/ext/NNlibCUDACUDNNExt/batchnorm.jl
@@ -186,9 +186,7 @@ function EnzymeCore.EnzymeRules.augmented_primal(config, func::EnzymeCore.Const{
     cache_running_var = nothing
 
     if !(typeof(y) <: EnzymeCore.Const)
-      if !(typeof(x) <: EnzymeCore.Const)
-         || !(typeof(g) <: EnzymeCore.Const)
-         || !(typeof(b) <: EnzymeCore.Const)
+      if !(typeof(x) <: EnzymeCore.Const) || !(typeof(g) <: EnzymeCore.Const) || !(typeof(b) <: EnzymeCore.Const)
 
         if EnzymeCore.EnzymeRules.overwritten(config)[3]
           cache_g = copy(g.val)
@@ -218,9 +216,7 @@ function EnzymeCore.EnzymeRules.reverse(config, func::EnzymeCore.Const{typeof(cu
     cache_g, cache_x, cache_running_mean, cache_running_var = cache
 
     if !(typeof(y) <: EnzymeCore.Const)
-      if !(typeof(x) <: EnzymeCore.Const)
-         || !(typeof(g) <: EnzymeCore.Const)
-         || !(typeof(b) <: EnzymeCore.Const)
+      if !(typeof(x) <: EnzymeCore.Const) || !(typeof(g) <: EnzymeCore.Const) || !(typeof(b) <: EnzymeCore.Const)
 
         if EnzymeCore.EnzymeRules.overwritten(config)[3]
           cache_g = g.val

diff --git a/src/NNlib.jl b/src/NNlib.jl
@@ -123,6 +123,4 @@ include("impl/depthwiseconv_im2col.jl")
 include("impl/pooling_direct.jl")
 include("deprecations.jl")
 
-include("enzyme.jl")
-
 end # module NNlib
diff --git a/src/enzyme.jl b/src/enzyme.jl