Trouble writing custom rule in Enzyme: AssertionError: !(overwritten[end])
#1242
Comments
|
What version of Enzyme are you on?, can you should the result of |
(@v1.9) pkg> status Enzyme
Status `C:\Users\gianl\.julia\environments\v1.9\Project.toml`
⌃ [7da242da] Enzyme v0.11.4
Info Packages marked with ⌃ have new versions available and may be upgradable.julia> versioninfo()
Julia Version 1.9.4
Commit 8e5136fa29 (2023-11-14 08:46 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Windows (x86_64-w64-mingw32)
CPU: 4 × Intel(R) Core(TM) i5-7600K CPU @ 3.80GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-14.0.6 (ORCJIT, skylake)
Threads: 1 on 4 virtual cores |
|
You're on a very old version of Enzyme, can you update? |
|
Yes, I will update and see if I still get an error. Edit: |
|
I've now tested on version 0.11.12 and the error persists: (@v1.9) pkg> status Enzyme
Status `C:\Users\gianl\.julia\environments\v1.9\Project.toml`
[7da242da] Enzyme v0.11.12 |
|
So I believe that this error can likely be removed (and is in fact an extraneous assertion). However, I'm still investigating. It would be really helpful to have a self-contained example (e.g. no external packages) if you're able to make a reproducer like that. |
|
The using Enzyme
using StaticArrays
import .EnzymeRules: forward, reverse, augmented_primal
using .EnzymeRules
const zero3 = @SArray [
0.0+0.0im 0.0+0.0im 0.0+0.0im
0.0+0.0im 0.0+0.0im 0.0+0.0im
0.0+0.0im 0.0+0.0im 0.0+0.0im
]
const eye3 = @SArray [
1.0+0.0im 0.0+0.0im 0.0+0.0im
0.0+0.0im 1.0+0.0im 0.0+0.0im
0.0+0.0im 0.0+0.0im 1.0+0.0im
]
const unitvec = ((1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1))
# U is filled with random special unitary matrices, but would make MWE too long
U = Array{SMatrix{3, 3, ComplexF64, 9}, 5}(undef, 4, 4, 4, 4, 4); fill!(U, eye3);
dU = similar(U); fill!(dU, zero3);
traceless_antihermitian(M::SMatrix{3,3,ComplexF64,9}) = 0.5*(M - M') - 1/6*tr(M - M')*eye3
tr(M::SMatrix{3,3,ComplexF64,9}) = M[1, 1] + M[2, 2] + M[3, 3]
@inline function smove(s::CartesianIndex{4}, μ, steps, lim)
newI = mod1.(s.I .+ steps .* unitvec[μ], lim)
return CartesianIndex(newI)
end
@generated function remultr(args::Vararg{T, N}) where {N,T}
quote
$(Expr(:meta, :inline))
tmp = *(args...)
real(tr(tmp))
end
end
@inline function plaquette(U, μ, ν, site)
Nμ = size(U)[1+μ]
Nν = size(U)[1+ν]
siteμ⁺ = smove(site, μ, 1, Nμ)
siteν⁺ = smove(site, ν, 1, Nν)
return remultr(U[μ,site], U[ν,siteμ⁺], U[μ,siteν⁺], U[ν,site])
end
function plaquette_sum(U::Array{SMatrix{3,3,ComplexF64,9}, 5})
p = 0.0
for site in CartesianIndices(size(U)[2:end])
for μ in 1:3
for ν in μ+1:4
p += plaquette(U, μ, ν, site)
end
end
end
return 6.0 * (6*prod(size(U)[2:end]) - 1/3*p)
end
### RULES FOR REMULTR ###
@inline function Base.circshift(shift::Integer, args::Vararg{T, N}) where {N,T}
j = mod1(shift, N)
ntuple(k -> args[k-j+ifelse(k>j,0,N)], Val(N))
end
function augmented_primal(config::ConfigWidth{1}, func::Const{typeof(remultr)},
::Type{<:Active}, args::Vararg{Active,N}) where {N}
argvals = ntuple(i -> args[i].val, Val(N))
if needs_primal(config)
primal = func.val(argvals...)
else
primal = nothing
end
if overwritten(config)[3]
tape = copy(argvals)
else
tape = nothing
end
return AugmentedReturn(primal, nothing, tape)
end
function EnzymeRules.reverse(config::ConfigWidth{1}, func::Const{typeof(remultr)},
dret::Active, tape, args::Vararg{Active,N}) where {N}
argvals = ntuple(i -> args[i].val, Val(N))
dargs = ntuple(Val(N)) do i
0.5traceless_antihermitian(*(circshift(i-1, argvals...)...))
end
return dargs
end
function Enzyme.gradient(::ReverseMode, f::typeof(remultr), args::Vararg{T,N}) where {N,T}
annots = ntuple(i -> Active(args[i]), Val(N))
der = autodiff(Reverse, f, Active, annots...)
return der
end
### RULES FOR PLAQUETTE_SUM ###
function Enzyme.gradient!(::ReverseMode, dU::Array{SMatrix{3,3,ComplexF64,9},5},
f::typeof(plaquette_sum), U::Array{SMatrix{3,3,ComplexF64,9},5})
autodiff(Reverse, f, Active, DuplicatedNoNeed(U, dU))
return nothing
end
# matrices = [eye3 for _ in 1:4] # should be any random special unitary matrices
# dm = Enzyme.gradient(Reverse, remultr, matrices...) # works as wanted
Enzyme.gradient!(Reverse, dU, plaquette_sum, U) # AssertionError: !(overwritten[end])I still get the error, but with slightly different stacktrace: |
using Enzyme
Enzyme.API.printall!(true)
import .EnzymeRules: forward, reverse, augmented_primal
using .EnzymeRules
function remultr(arg)
real(arg)
end
function augmented_primal(config::ConfigWidth{1}, func::Const{typeof(remultr)},
::Type{<:Active}, args::Vararg{Active,N}) where {N}
return AugmentedReturn(func.val(args[1].val), nothing, nothing)
end
function EnzymeRules.reverse(config::ConfigWidth{1}, func::Const{typeof(remultr)},
dret::Active, tape, args::Vararg{Active,N}) where {N}
dargs = ntuple(Val(N)) do i
0.5
end
return dargs
end
# U is filled with random special unitary matrices, but would make MWE too long
U = Array{Complex{Float64}, 2}(undef, 4, 4)
dU = similar(U)
function plaquette_sum(U)
p = 0.0
for site in CartesianIndices(size(U)[2:end])
p += remultr(@inbounds U[1,site]) # , U[1,site])
end
return p
end
@show collect(CartesianIndices(size(U)[2:end]))
autodiff(Reverse, plaquette_sum, Active, DuplicatedNoNeed(U, dU)) |
|
Okay, I've deduced this issue is due to some lacking lifetime information. The PR I linked above is a workaround for 1.10. The full actual solution is for JuliaLang/julia#53095 to be merged and backported to whatever LLVM version you're using. |
|
Wow, I'm glad you were able to figure out the root cause despite my contrived example! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
(Transferred issue from Julia discourse: https://discourse.julialang.org/t/trouble-writing-custom-rule-in-enzyme-assertionerror-overwritten-end/108939)
Hi,
I am currently trying to use Enzyme to get the derivative of a function that takes a multidimensional array of 3x3 matrices and outputs a scalar. The caveat is that these matrices are elements of the Lie-group of special unitary matrices SU(3) and therefore the derivatives w.r.t. each element in the array should be in the corresponding algebra (traceless anti-Hermitian matrices).
The function in question is:
I've written a reverse-mode rule for
remultr(which I will define later), directly following the example in the docs, that works well enough (and without error) for this example. I also overloaded thegradient!function forplaquette_sum, that results in an error:Since writing the discourse post, I found that whether I get the error on
gradient!(...)depends on how I define theremultrfunction.When I use
I get no error, but if I use:
the error appears. The reason the
@generatedsecond definition exists, is because I wrote a custom matrix-multiplication routing using LoopVectorization.jl 's@turboin a package I want to put this whole thing into. Ultimately, I should be able use a definition ofremultrsimilar to the first and be fine, but I still wanted to post this issue here in case it could be useful. I hope this is enough to recreate the error.P.S.: Here is the complete error message
The text was updated successfully, but these errors were encountered: