Small inference test
- 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
Assessment
This issue has not been assessed yet.