EnzymeAD / EnzymeAD/Enzyme.jl

Small inference test

Open
#2,152 5 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

Pasted here for ease.

What I'm using locally for some testing @vchuravy

```
module ReverseRules

using Enzyme
using Enzyme: EnzymeRules
using LinearAlgebra
using Test

f(x) = x^2

function f_ip(x)
x[1] *= x[1]
return nothing
end

import .EnzymeRules: augmented_primal, reverse, Annotation, has_rrule_from_sig
using .EnzymeRules

q(x) = x^2
function augmented_primal(config::RevConfigWidth{1}, func::Const{typeof(q)}, ::Type{<:Active}, x::Active)
tape = (Ref(2.0), Ref(3.4))
if needs_primal(config)
return AugmentedReturn(func.val(x.val), nothing, tape)
else
return AugmentedReturn(nothing, nothing, tape)
end
end

function reverse(config::RevConfigWidth{1}, ::Const{typeof(q)}, dret::Active, tape, x::Active)
@test tape[1][] == 2.0
@test tape[2][] == 3.4
if needs_primal(config)
return (10+2*x.val*dret.val,)
else
return (100+2*x.val*dret.val,)
end
end

using Profile
@profile Enzyme.autodiff(Enzyme.Reverse, q, Active(2.0))
using PProf
pprof(; webhost="")
while true

end

using SnoopCompileCore

tinf = @snoop_inference Enzyme.autodiff(Enzyme.Reverse, q, Active(2.0))
using SnoopCompile
using AbstractTrees
out=open("dat.txt", "w")
print_tree(out, tinf, maxdepth=100)
close(out)

out=open("inv.txt", "w")
@show invalidation_trees(tinf)
println(out, invalidation_trees(tinf))
close(out)

#@testset "Byref Tape" begin
# @test Enzyme.autodiff(Enzyme.Reverse, q, Active(2.0))[1][1] ≈ 104.0
#end

end # ReverseRules
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.