# Checkpointing.jl does not work in the latest version of Enzyme.jl

**URL:** <https://discourse.julialang.org/t/checkpointing-jl-does-not-work-in-the-latest-version-of-enzyme-jl/118660>\
**Category:** Numerics\
**Tags:** question, enzyme\
**Created:** [August 27, 2024, 12:54pm UTC](https://discourse.julialang.org/t/checkpointing-jl-does-not-work-in-the-latest-version-of-enzyme-jl/118660 "2024-08-27T12:54:10Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![ToPo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/topo/32/205868_2.png) [@ToPo](https://discourse.julialang.org/u/ToPo)\
**Post date:** [August 27, 2024, 12:54pm UTC](https://discourse.julialang.org/t/checkpointing-jl-does-not-work-in-the-latest-version-of-enzyme-jl/118660/1 "2024-08-27T12:54:10Z")

</div>

Recently, I have been using the checkpointing method to reduce the memory required for automatic differentiation.  
Therefore, we use Enzyme.jl and Checkpointing.jl.  
However, when I use the latest version of Enzyme.jl, **v0.12.33** , an error occurs.  
The version of Checkpointing.jl is **v0.9.4**.

The calculations work fine until **v0.12.27** of Enzyme.jl.  
**What change in Enzyme.jl is causing this?**

Below is a sample program and execution results of Checkpointing.jl.

```julia
# Explicit 1D heat equation
using Checkpointing
using Enzyme
using Plots

mutable struct Heat
    Tnext::Vector{Float64}
    Tlast::Vector{Float64}
    n::Int
    λ::Float64
    tsteps::Int
    step::Int64
end

function advance(heat)
    next = heat.Tnext
    last = heat.Tlast
    λ = heat.λ
    n = heat.n
    for i in 2:(n-1)
        next[i] = last[i] + λ*(last[i-1]-2*last[i]+last[i+1])
    end
    return nothing
end

function sumheat(heat::Heat, chkpscheme::Scheme, tsteps::Int64)
    # AD: Create shadow copy for derivatives
    @checkpoint_struct chkpscheme heat for i in 1:tsteps
        heat.Tlast .= heat.Tnext
        advance(heat)
    end
    return reduce(+, heat.Tnext)
end

function heat(scheme::Scheme, tsteps::Int)
    n = 100
    Δx=0.1
    Δt=0.001
    # Select μ such that λ ≤ 0.5 for stability with μ = (λ*Δt)/Δx^2
    λ = 0.5

    # Create object from struct. tsteps is not needed for a for-loop
    heat = Heat(zeros(n), zeros(n), n, λ, tsteps,1)
    # Shadow copy for Enzyme
    dheat = Heat(zeros(n), zeros(n), n, λ, tsteps,1)

    # Boundary conditions
    heat.Tnext[1] = 20.0
    heat.Tnext[end] = 0

    # Compute gradient
    autodiff(Enzyme.ReverseWithPrimal, sumheat, Duplicated(heat, dheat), Const(scheme), Const(tsteps))

    return heat.Tnext, dheat.Tnext[2:end-1]
end
tsteps = 5
T, dT = heat(Revolve{Heat}(tsteps,4; verbose=1), tsteps)
# Plot function values
plot(T)
# Plot gradient with respect of sum(T[end]) with respect to T[1].
plot(dT)

```

```julia
[ Info: Revolve: Number of checkpoints: 4
[ Info: Revolve: Number of steps: 5
[ Info: Prediction:
[ Info: Forward steps : 4
[ Info: Overhead factor : 0.8
[ Info: Size per checkpoint: 1.688 KiB
[ Info: Revolve: First Uturn
[ Info: Size of total storage: 6.828 KiB
ERROR: Function argument passed to autodiff cannot be proven readonly.
If the the function argument cannot contain derivative data, instead call autodiff(Mode, Const(f), ...)
See https://enzyme.mit.edu/index.fcgi/julia/stable/faq/#Activity-of-temporary-storage for more information.
The potentially writing call is %67 = call nonnull "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* ({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)*, {} addrspace(10)*, ...) @julia.call({} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32)* noundef nonnull @ijl_apply_generic, {} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 4829008240 to {}*) to {} addrspace(10)*), {} addrspace(10)* nonnull %52, {} addrspace(10)* addrspacecast ({}* inttoptr (i64 4367827296 to {}*) to {} addrspace(10)*)) #22, !dbg !202, using %52 = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %51 unordered, align 8, !dbg !202, !tbaa !34, !alias.scope !38, !noalias !41
Stacktrace:
 [1] #16
   @ ~/.julia/packages/Checkpointing/CEFSa/src/Checkpointing.jl:182
 [2] #16
   @ ~/.julia/packages/Checkpointing/CEFSa/src/Checkpointing.jl:0

Stacktrace:
  [1] #16
    @ ~/.julia/packages/Checkpointing/CEFSa/src/Checkpointing.jl:182 [inlined]
  [2] diffejulia__16_7917_inner_1wrap
    @ ~/.julia/packages/Checkpointing/CEFSa/src/Checkpointing.jl:0
  [3] macro expansion
    @ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:7099 [inlined]
  [4] enzyme_call
    @ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:6708 [inlined]
  [5] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:6585 [inlined]
  [6] autodiff
    @ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:320 [inlined]
  [7] autodiff
    @ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:348 [inlined]
  [8] autodiff
    @ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:329 [inlined]
  [9] rev_checkpoint_struct_for
    @ ~/.julia/packages/Checkpointing/CEFSa/src/Schemes/Revolve.jl:471
 [10] macro expansion
    @ ~/.julia/packages/Checkpointing/CEFSa/src/Checkpointing.jl:176 [inlined]
 [11] sumheat
    @ ~/my_program/GitHub/JuliaOptOS/tests/Enzyme/check_pointing/check_pointing_1d_heat_test.jl:31 [inlined]
 [12] diffejulia_sumheat_7819wrap
    @ ~/my_program/GitHub/JuliaOptOS/tests/Enzyme/check_pointing/check_pointing_1d_heat_test.jl:0
 [13] macro expansion
    @ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:7099 [inlined]
 [14] enzyme_call
    @ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:6708 [inlined]
 [15] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/YWQiS/src/compiler.jl:6585 [inlined]
 [16] autodiff
    @ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:320 [inlined]
 [17] autodiff
    @ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:348 [inlined]
 [18] autodiff
    @ ~/.julia/packages/Enzyme/YWQiS/src/Enzyme.jl:329 [inlined]
 [19] heat(scheme::Revolve{Heat}, tsteps::Int64)
    @ Main ~/my_program/GitHub/JuliaOptOS/tests/Enzyme/check_pointing/check_pointing_1d_heat_test.jl:55
 [20] top-level scope
    @ ~/my_program/GitHub/JuliaOptOS/tests/Enzyme/check_pointing/check_pointing_1d_heat_test.jl:60

```

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [August 27, 2024, 3:08pm UTC](https://discourse.julialang.org/t/checkpointing-jl-does-not-work-in-the-latest-version-of-enzyme-jl/118660/2 "2024-08-27T15:08:01Z")

</div>

@Michel_Schanen

---

<div class="post-metadata">

**Author:** ![ToPo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/topo/32/205868_2.png) [@ToPo](https://discourse.julialang.org/u/ToPo)\
**Post date:** [September 9, 2024, 6:59am UTC](https://discourse.julialang.org/t/checkpointing-jl-does-not-work-in-the-latest-version-of-enzyme-jl/118660/3 "2024-09-09T06:59:46Z")

</div>

Thank you.  
The version of Checkpointing.jl was updated.  
This confirms that the calculation can be performed with the latest version v0.12.36 of Enzyme.jl and v0.9.5 of Checkpointing.jl.