Try to test Enzyme on 1.10 (#270)

* Try to test Enzyme on 1.10

* Remove nonexistent test

* Fix forward mode returns

Author: kshyatt

Project.toml

@@ -20,10 +20,10 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2"
+cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
-cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [extensions]
 TensorOperationsBumperExt = "Bumper"
@@ -42,7 +42,7 @@ CUDACore = "6"
 ChainRulesCore = "1"
 ChainRulesTestUtils = "1"
 DynamicPolynomials = "0.5, 0.6"
-Enzyme = "0.13.115"
+Enzyme = "0.13.146"
 EnzymeTestUtils = "0.2"
 JLArrays = "0.3"
 LRUCache = "1"
@@ -74,6 +74,7 @@ cuRAND = "20fd9a0b-12d5-4c2f-a8af-7c34e9e60431"
 DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl

@@ -164,7 +164,15 @@ function EnzymeRules.forward(
         end
     end
     TensorOperations.tensorcontract!(C_dC.val, A_dA.val, pA, conjA, B_dB.val, pB, conjB, pAB, α, β, ba...)
-    return C_dC
+    if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config)
+        return C_dC
+    elseif EnzymeRules.needs_primal(config)
+        return C_dC.val
+    elseif EnzymeRules.needs_shadow(config)
+        return C_dC.dval
+    else
+        return nothing
+    end
 end
 
 function EnzymeRules.augmented_primal(
@@ -275,7 +283,15 @@ function EnzymeRules.forward(
         end
     end
     TensorOperations.tensoradd!(C_dC.val, A_dA.val, pA, conjA, α, β, ba...)
-    return C_dC
+    if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config)
+        return C_dC
+    elseif EnzymeRules.needs_primal(config)
+        return C_dC.val
+    elseif EnzymeRules.needs_shadow(config)
+        return C_dC.dval
+    else
+        return nothing
+    end
 end
 
 function EnzymeRules.augmented_primal(
@@ -389,7 +405,15 @@ function EnzymeRules.forward(
     end
     # D = α * tr(A) + β * C
     TensorOperations.tensortrace!(C_dC.val, A_dA.val, p, q, conjA, α, β, ba...)
-    return C_dC
+    if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config)
+        return C_dC
+    elseif EnzymeRules.needs_primal(config)
+        return C_dC.val
+    elseif EnzymeRules.needs_shadow(config)
+        return C_dC.dval
+    else
+        return nothing
+    end
 end
 
 end

test/runtests.jl

@@ -39,13 +39,10 @@ if !is_buildkite
     @testset "mooncake" verbose = false begin
         include("mooncake.jl")
     end
-    # mystery segfault on 1.10 for now
-    @static if VERSION >= v"1.11.0"
-        is_apple_ci = Sys.isapple() && get(ENV, "CI", "false") == "true"
-        if !is_apple_ci
-            @testset "enzyme" verbose = false begin
-                include("enzyme.jl")
-            end
+    is_apple_ci = Sys.isapple() && get(ENV, "CI", "false") == "true"
+    if !is_apple_ci
+        @testset "enzyme" verbose = false begin
+            include("enzyme.jl")
         end
     end
 end