attempt fix

Author: wsmoses

src/NNlib.jl

@@ -123,10 +123,6 @@ include("impl/depthwiseconv_im2col.jl")
 include("impl/pooling_direct.jl")
 include("deprecations.jl")
 
-@init @static if !isdefined(Base, :get_extension)
-    @require EnzymeCore="f151be2c-9106-41f4-ab19-57ee4f262869" begin
-        include("../ext/NNlibEnzymeCoreExt/NNlibEnzymeCoresExt.jl")
-    end
-end
+include("enzyme.jl")
 
 end # module NNlib

test/conv.jl

@@ -870,7 +870,7 @@ end
   w = rand(rng, repeat([3], spatial_rank)..., 3, 3)
   cdims = DenseConvDims(x, w)
   gradtest((x, w) -> conv(x, w, cdims), x, w)
-  gradtest((x, w) -> sum(conv(x, w, cdims)), x, w)  # https://github.com/FluxML/Flux.jl/issues/1055
+  gradtest((x, w) -> sum(conv(x, w, cdims)), x, w; check_enzyme_rule=true)  # https://github.com/FluxML/Flux.jl/issues/1055
 
   y = conv(x, w, cdims)
   gradtest((y, w) -> ∇conv_data(y, w, cdims), y, w)

test/test_utils.jl

@@ -12,14 +12,30 @@ Applies also `ChainRulesTestUtils.test_rrule` if the rrule for `f` is explicitly
 """
 function gradtest(
     f, xs...; atol = 1e-6, rtol = 1e-6, fkwargs = NamedTuple(),
-    check_rrule = false, fdm = :central, check_broadcast = false,
+    check_rrule = false, check_enzyme_rrule = false, fdm = :central, check_broadcast = false,
     skip = false, broken = false,
 )
     # TODO: revamp when https://github.com/JuliaDiff/ChainRulesTestUtils.jl/pull/166
     # is merged
     if check_rrule
         test_rrule(f, xs...; fkwargs = fkwargs)
     end
+    if check_enzyme_rrule
+        if len(xs) == 2
+            for Tret in (Const, Active),
+                Tx in (Const, Duplicated, BatchDuplicated),
+                Ty in (Const, Duplicated, BatchDuplicated)
+
+                are_activities_compatible(Tret, Tx, Ty) || continue
+
+                test_reverse(fun, Tret, (xs[1], Tx), (ys[1], Ty); atol, rtol)
+            end
+        else
+            throw(AssertionError("Unsupported arg count for testing"))
+        end
+
+        EnzymeTestUtils.test_rrule(f, xs...; fkwargs = fkwargs)
+    end
 
     if check_broadcast
         length(fkwargs) > 0 && @warn("CHECK_BROADCAST: dropping keywords args")