Merge branch 'main' into recursive_accumulate

EnzymeAD · Nov 18, 2024 · 05efebf · 05efebf
2 parents 530c780 + ba4c22a
commit 05efebf
Show file tree

Hide file tree

Showing 29 changed files with 2,206 additions and 788 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -199,7 +199,7 @@ jobs:
       matrix:
         version:
           - '1.10'
-          - ~1.11.0-0
+          - '1.11'
           - 'nightly'
         os:
           - ubuntu-latest
@@ -264,6 +264,7 @@ jobs:
           - ubuntu-latest
         test:
           - DynamicExpressions
+          - Bijectors
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -273,8 +274,8 @@ jobs:
       - uses: julia-actions/julia-buildpkg@v1
       - name: "Run tests"
         run: |
-            julia --color=yes --project=test/integration -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in (".", "lib/EnzymeCore")]); Pkg.instantiate()'
-            julia --color=yes --project=test/integration --threads=auto --check-bounds=yes test/integration/${{ matrix.test }}.jl
+            julia --color=yes --project=test/integration/${{ matrix.test }} -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in (".", "lib/EnzymeCore")]); Pkg.instantiate()'
+            julia --color=yes --project=test/integration/${{ matrix.test }} --threads=auto --check-bounds=yes test/integration/${{ matrix.test }}/runtests.jl
         shell: bash
   docs:
     name: Documentation

diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Enzyme"
 uuid = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>"]
-version = "0.13.13"
+version = "0.13.15"
 
 [deps]
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
@@ -35,8 +35,8 @@ EnzymeStaticArraysExt = "StaticArrays"
 BFloat16s = "0.2, 0.3, 0.4, 0.5"
 CEnum = "0.4, 0.5"
 ChainRulesCore = "1"
-EnzymeCore = "0.8.4, 0.8.5"
-Enzyme_jll = "0.0.158"
+EnzymeCore = "0.8.6"
+Enzyme_jll = "0.0.165"
 GPUCompiler = "0.21, 0.22, 0.23, 0.24, 0.25, 0.26, 0.27, 1"
 LLVM = "6.1, 7, 8, 9"
 LogExpFunctions = "0.3"

diff --git a/lib/EnzymeCore/Project.toml b/lib/EnzymeCore/Project.toml
@@ -1,7 +1,7 @@
 name = "EnzymeCore"
 uuid = "f151be2c-9106-41f4-ab19-57ee4f262869"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>"]
-version = "0.8.5"
+version = "0.8.6"
 
 [compat]
 Adapt = "3, 4"

diff --git a/lib/EnzymeCore/src/EnzymeCore.jl b/lib/EnzymeCore/src/EnzymeCore.jl
@@ -332,6 +332,7 @@ const ReverseHolomorphicWithPrimal = ReverseMode{true,false,DefaultABI, true, fa
 @inline set_err_if_func_written(::ReverseMode{ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten}) where {ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten} = ReverseMode{ReturnPrimal,RuntimeActivity,ABI,Holomorphic,true}()
 @inline clear_err_if_func_written(::ReverseMode{ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten}) where {ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten} = ReverseMode{ReturnPrimal,RuntimeActivity,ABI,Holomorphic,false}()
 
+@inline set_abi(::Type{ReverseMode{ReturnPrimal,RuntimeActivity,OldABI,Holomorphic,ErrIfFuncWritten}}, ::Type{NewABI}) where {ReturnPrimal,RuntimeActivity,OldABI,Holomorphic,ErrIfFuncWritten,NewABI<:ABI} = ReverseMode{ReturnPrimal,RuntimeActivity,NewABI,Holomorphic,ErrIfFuncWritten}
 @inline set_abi(::ReverseMode{ReturnPrimal,RuntimeActivity,OldABI,Holomorphic,ErrIfFuncWritten}, ::Type{NewABI}) where {ReturnPrimal,RuntimeActivity,OldABI,Holomorphic,ErrIfFuncWritten,NewABI<:ABI} = ReverseMode{ReturnPrimal,RuntimeActivity,NewABI,Holomorphic,ErrIfFuncWritten}()
 
 @inline set_runtime_activity(::ReverseMode{ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten}) where {ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten} = ReverseMode{ReturnPrimal,true,ABI,Holomorphic,ErrIfFuncWritten}()
@@ -483,6 +484,7 @@ const ForwardWithPrimal = ForwardMode{true, DefaultABI, false, false}()
 @inline set_err_if_func_written(::ForwardMode{ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity}) where {ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity} = ForwardMode{ReturnPrimal,ABI,true,RuntimeActivity}()
 @inline clear_err_if_func_written(::ForwardMode{ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity}) where {ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity} = ForwardMode{ReturnPrimal,ABI,false,RuntimeActivity}()
 
+@inline set_abi(::Type{ForwardMode{ReturnPrimal,OldABI,ErrIfFuncWritten,RuntimeActivity}}, ::Type{NewABI}) where {ReturnPrimal,OldABI,ErrIfFuncWritten,RuntimeActivity,NewABI<:ABI} = ForwardMode{ReturnPrimal,NewABI,ErrIfFuncWritten,RuntimeActivity}
 @inline set_abi(::ForwardMode{ReturnPrimal,OldABI,ErrIfFuncWritten,RuntimeActivity}, ::Type{NewABI}) where {ReturnPrimal,OldABI,ErrIfFuncWritten,RuntimeActivity,NewABI<:ABI} = ForwardMode{ReturnPrimal,NewABI,ErrIfFuncWritten,RuntimeActivity}()
 
 @inline set_runtime_activity(::ForwardMode{ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity}) where {ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity} = ForwardMode{ReturnPrimal,ABI,ErrIfFuncWritten,true}()
@@ -667,4 +669,60 @@ Return a new mode with its [`ABI`](@ref) set to the chosen type.
 """
 function set_abi end
 
+
+"""
+    Primitive Type usable within Reactant. See Reactant.jl for more information.
+"""
+@static if isdefined(Core, :BFloat16)
+    const ReactantPrimitive = Union{
+        Bool,
+        Int8,
+        UInt8,
+        Int16,
+        UInt16,
+        Int32,
+        UInt32,
+        Int64,
+        UInt64,
+        Float16,
+        Core.BFloat16,
+        Float32,
+        Float64,
+        Complex{Float32},
+        Complex{Float64},
+    }
+else
+    const ReactantPrimitive = Union{
+        Bool,
+        Int8,
+        UInt8,
+        Int16,
+        UInt16,
+        Int32,
+        UInt32,
+        Int64,
+        UInt64,
+        Float16,
+        Float32,
+        Float64,
+        Complex{Float32},
+        Complex{Float64},
+    }
+end
+
+"""
+    Abstract Reactant Array type. See Reactant.jl for more information
+"""
+abstract type RArray{T<:ReactantPrimitive,N} <: AbstractArray{T,N} end
+@inline Base.eltype(::RArray{T}) where T = T
+@inline Base.eltype(::Type{<:RArray{T}}) where T = T
+
+"""
+    Abstract Reactant Number type. See Reactant.jl for more information
+"""
+abstract type RNumber{T<:ReactantPrimitive} <: Number end
+@inline Base.eltype(::RNumber{T}) where T = T
+@inline Base.eltype(::Type{<:RNumber{T}}) where T = T
+
+
 end # module EnzymeCore
diff --git a/lib/EnzymeCore/test/runtests.jl b/lib/EnzymeCore/test/runtests.jl
@@ -3,31 +3,31 @@ using EnzymeCore
 
 @testset verbose = true "EnzymeCore" begin
     @testset "WithPrimal" begin
-        @test WithPrimal(Reverse) === ReverseWithPrimal
-        @test NoPrimal(Reverse) === Reverse
-        @test WithPrimal(ReverseWithPrimal) === ReverseWithPrimal
-        @test NoPrimal(ReverseWithPrimal) === Reverse
+        @test EnzymeCore.WithPrimal(Reverse) === ReverseWithPrimal
+        @test EnzymeCore.NoPrimal(Reverse) === Reverse
+        @test EnzymeCore.WithPrimal(ReverseWithPrimal) === ReverseWithPrimal
+        @test EnzymeCore.NoPrimal(ReverseWithPrimal) === Reverse
 
-        @test WithPrimal(set_runtime_activity(Reverse)) === set_runtime_activity(ReverseWithPrimal)
+        @test EnzymeCore.WithPrimal(EnzymeCore.set_runtime_activity(Reverse)) === EnzymeCore.set_runtime_activity(ReverseWithPrimal)
 
-        @test WithPrimal(Forward) === ForwardWithPrimal
-        @test NoPrimal(Forward) === Forward
-        @test WithPrimal(ForwardWithPrimal) === ForwardWithPrimal
-        @test NoPrimal(ForwardWithPrimal) === Forward
+        @test EnzymeCore.WithPrimal(Forward) === ForwardWithPrimal
+        @test EnzymeCore.NoPrimal(Forward) === Forward
+        @test EnzymeCore.WithPrimal(ForwardWithPrimal) === ForwardWithPrimal
+        @test EnzymeCore.NoPrimal(ForwardWithPrimal) === Forward
 
-        @test WithPrimal(ReverseSplitNoPrimal) === ReverseSplitWithPrimal
-        @test NoPrimal(ReverseSplitNoPrimal) === ReverseSplitNoPrimal
-        @test WithPrimal(ReverseSplitWithPrimal) === ReverseSplitWithPrimal
-        @test NoPrimal(ReverseSplitWithPrimal) === ReverseSplitNoPrimal
+        @test EnzymeCore.WithPrimal(ReverseSplitNoPrimal) === ReverseSplitWithPrimal
+        @test EnzymeCore.NoPrimal(ReverseSplitNoPrimal) === ReverseSplitNoPrimal
+        @test EnzymeCore.WithPrimal(ReverseSplitWithPrimal) === ReverseSplitWithPrimal
+        @test EnzymeCore.NoPrimal(ReverseSplitWithPrimal) === ReverseSplitNoPrimal
     end
 
     @testset "needs_primal" begin
-        @test needs_primal(Reverse) === false
-        @test needs_primal(ReverseWithPrimal) === true
-        @test needs_primal(Forward) === false
-        @test needs_primal(ForwardWithPrimal) === true
-        @test needs_primal(ReverseSplitNoPrimal) === false
-        @test needs_primal(ReverseSplitWithPrimal) === true
+        @test EnzymeCore.needs_primal(Reverse) === false
+        @test EnzymeCore.needs_primal(ReverseWithPrimal) === true
+        @test EnzymeCore.needs_primal(Forward) === false
+        @test EnzymeCore.needs_primal(ForwardWithPrimal) === true
+        @test EnzymeCore.needs_primal(ReverseSplitNoPrimal) === false
+        @test EnzymeCore.needs_primal(ReverseSplitWithPrimal) === true
     end
 
     @testset "Miscellaneous" begin