Utilities for splitting and unsplitting mode objects

lib/EnzymeCore/src/EnzymeCore.jl

@@ -384,7 +384,7 @@ Subtype of [`Mode`](@ref) for split reverse mode differentiation, to use in [`au
     - [`set_abi`](@ref)
     - [`ReverseSplitModified`](@ref), [`ReverseSplitWidth`](@ref)
 """
-struct ReverseModeSplit{ReturnPrimal,ReturnShadow,Width,RuntimeActivity,ModifiedBetween,ABI,Holomorphic,ErrIfFuncWritten,ShadowInit} <: Mode{ABI, ErrIfFuncWritten,RuntimeActivity} end
+struct ReverseModeSplit{ReturnPrimal,ReturnShadow,RuntimeActivity,Width,ModifiedBetween,ABI,Holomorphic,ErrIfFuncWritten,ShadowInit} <: Mode{ABI, ErrIfFuncWritten,RuntimeActivity} end
 
 """
     const ReverseSplitNoPrimal
@@ -578,4 +578,88 @@ Return a new mode with its [`ABI`](@ref) set to the chosen type.
 """
 function set_abi end
 
+"""
+    split_mode(::ReverseMode)
+
+Turn a [`ReverseMode`](@ref) object into a [`ReverseModeSplit`](@ref) object while preserving as many of the settings as possible.
+
+This function acts as the identity on a [`ReverseModeSplit`](@ref).
+
+See also [`unsplit_mode`](@ref).
+"""
+function split_mode(::ReverseMode{
+    ReturnPrimal,
+    RuntimeActivity,
+    ABI,
+    Holomorphic,
+    ErrIfFuncWritten
+}) where {
+    ReturnPrimal,
+    RuntimeActivity,
+    ABI,
+    Holomorphic,
+    ErrIfFuncWritten
+}
+    ReturnShadow = true
+    Width = 0
+    ModifiedBetween = true
+    ShadowInit = false
+    mode_split = ReverseModeSplit{
+        ReturnPrimal,
+        ReturnShadow,
+        RuntimeActivity,
+        Width,
+        ModifiedBetween,
+        ABI,
+        Holomorphic,
+        ErrIfFuncWritten,
+        ShadowInit
+    }()
+    return mode_split
+end
+
+split_mode(mode::ReverseModeSplit) = mode
+
+"""
+    unsplit_mode(::ReverseMode)
+
+Turn a [`ReverseModeSplit`](@ref) object into a [`ReverseMode`](@ref) object while preserving as many of the settings as possible.
+
+This function acts as the identity on a [`ReverseMode`](@ref).
+
+See also [`split_mode`](@ref).
+"""
+function unsplit_mode(::ReverseModeSplit{
+    ReturnPrimal,
+    ReturnShadow,
+    RuntimeActivity,
+    Width,
+    ModifiedBetween,
+    ABI,
+    Holomorphic,
+    ErrIfFuncWritten,
+    ShadowInit
+}) where {
+    ReturnPrimal,
+    ReturnShadow,
+    RuntimeActivity,
+    Width,
+    ModifiedBetween,
+    ABI,
+    Holomorphic,
+    ErrIfFuncWritten,
+    ShadowInit
+}
+    mode_unsplit = ReverseMode{
+        ReturnPrimal,
+        RuntimeActivity,
+        ABI,
+        Holomorphic,
+        ErrIfFuncWritten
+    }()
+    return mode_unsplit
+end
+
+unsplit_mode(mode::ReverseMode) = mode
+
 end # module EnzymeCore

lib/EnzymeCore/test/mode_modification.jl

@@ -0,0 +1,15 @@
+using EnzymeCore
+using EnzymeCore: split_mode, unsplit_mode
+using Test
+
+@testset "Split / unsplit mode" begin
+    @test split_mode(Reverse) == ReverseSplitNoPrimal
+    @test split_mode(ReverseWithPrimal) == ReverseSplitWithPrimal
+    @test split_mode(ReverseSplitNoPrimal) == ReverseSplitNoPrimal
+    @test split_mode(ReverseSplitWithPrimal) == ReverseSplitWithPrimal
+
+    @test unsplit_mode(Reverse) == Reverse
+    @test unsplit_mode(ReverseWithPrimal) == ReverseWithPrimal
+    @test unsplit_mode(ReverseSplitNoPrimal) == Reverse
+    @test unsplit_mode(ReverseSplitWithPrimal) == ReverseWithPrimal
+end

lib/EnzymeCore/test/runtests.jl

@@ -25,4 +25,6 @@ y = @view data[3:end]
 @test_throws ErrorException Duplicated(d, y)
 
 @test_throws ErrorException Active(data)
-@test_throws ErrorException Active(d)
+@test_throws ErrorException Active(d)
+
+include("mode_modification.jl")