Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Import needs_primal from Reactant.jl #2021

Merged
merged 8 commits into from
Oct 30, 2024
+44 −21
Merged

Import needs_primal from Reactant.jl #2021

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
3415d5d
Implement `needs_primal`
mofeing Oct 28, 2024
efd2925
Test `needs_primal`
mofeing Oct 28, 2024
4ebf7a8
Export `needs_primal`
mofeing Oct 29, 2024
7b6e394
Bump versions
mofeing Oct 29, 2024
2eaaf0e
Move `needs_primal` tests EnzymeCore
mofeing Oct 30, 2024
cede9c7
Use tilde version specifier to also support v0.8.4
mofeing Oct 30, 2024
9597c03
Move `WithPrimal`, `NoPrimal` tests to EnzymeCore
mofeing Oct 30, 2024
053f627
Merge branch 'main' into needs_primal
wsmoses Oct 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions lib/EnzymeCore/src/EnzymeCore.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -351,6 +351,15 @@ Return a new mode which excludes the primal value.
"""
@inline NoPrimal(::ReverseMode{ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten}) where {ReturnPrimal,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten} = ReverseMode{false,RuntimeActivity,ABI,Holomorphic,ErrIfFuncWritten}()

"""
needs_primal(::Mode)
needs_primal(::Type{Mode})

Returns `true` if the mode needs the primal value, otherwise `false`.
"""
@inline needs_primal(::ReverseMode{ReturnPrimal}) where {ReturnPrimal} = ReturnPrimal
@inline needs_primal(::Type{<:ReverseMode{ReturnPrimal}}) where {ReturnPrimal} = ReturnPrimal

"""
struct ReverseModeSplit{
ReturnPrimal,
Expand Down Expand Up @@ -424,6 +433,8 @@ Return a new instance of [`ReverseModeSplit`](@ref) mode where `Width` is set to
@inline WithPrimal(::ReverseModeSplit{ReturnPrimal,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit}) where {ReturnPrimal,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit} = ReverseModeSplit{true,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit}()
@inline NoPrimal(::ReverseModeSplit{ReturnPrimal,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit}) where {ReturnPrimal,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit} = ReverseModeSplit{false,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI, Holomorphic, ErrIfFuncWritten, ShadowInit}()

@inline needs_primal(::ReverseModeSplit{ReturnPrimal}) where {ReturnPrimal} = ReturnPrimal
@inline needs_primal(::Type{<:ReverseModeSplit{ReturnPrimal}}) where {ReturnPrimal} = ReturnPrimal

"""
struct ForwardMode{
Expand Down Expand Up @@ -480,6 +491,8 @@ const ForwardWithPrimal = ForwardMode{true, DefaultABI, false, false}()
@inline WithPrimal(::ForwardMode{ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity}) where {ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity} = ForwardMode{true,ABI,ErrIfFuncWritten,RuntimeActivity}()
@inline NoPrimal(::ForwardMode{ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity}) where {ReturnPrimal,ABI,ErrIfFuncWritten,RuntimeActivity} = ForwardMode{false,ABI,ErrIfFuncWritten,RuntimeActivity}()

@inline needs_primal(::ForwardMode{ReturnPrimal}) where {ReturnPrimal} = ReturnPrimal
@inline needs_primal(::Type{<:ForwardMode{ReturnPrimal}}) where {ReturnPrimal} = ReturnPrimal

function autodiff end
function autodiff_deferred end
Expand Down
9 changes: 9 additions & 0 deletions test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3634,6 +3634,15 @@ end
@test NoPrimal(ReverseSplitWithPrimal) === ReverseSplitNoPrimal
end

@testset "needs_primal" begin
@test needs_primal(Reverse) === false
mofeing marked this conversation as resolved.
Show resolved Hide resolved
@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
end

# TEST EXTENSIONS
using SpecialFunctions
@testset "SpecialFunctions ext" begin
Expand Down
Loading