EnzymeAD / EnzymeAD/Enzyme-JAX
autodiff through while loops with state-dependent break
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```julia
using Reactant
using Enzyme
mutable struct Model{X, T}
x::X
t::T
end
function step!(m::Model)
dt = 1 / sum(m.x)
m.t += dt
m.x .+= dt .* m.x
return nothing
end
function loss(m::Model, T)
Reactant.@trace checkpointing=false mincut=true track_numbers=false while (m.t < T)
step!(m)
end
return sum(m.x)
end
function grad_loss(m::Model, dm::Model, T)
Enzyme.autodiff(
Reverse, loss, Active,
Duplicated(m, dm), Const(T)
)
return dm
end
m = Model(Reactant.to_rarray([1.0, 1.0]), ConcreteRNumber(0.0))
T = ConcreteRNumber(10)
cl = Reactant.@compile raise=true sync=true loss(m, T)
@info cl(m, ConcreteRNumber(3))
```
```julia
[ Info: ConcretePJRTNumber{Float64, 1}(32.0)
```
```julia
dm = Enzyme.make_zero(m)
cgl = Reactant.@compile raise=true raise_first=true grad_loss(m, dm, T)
@info cgl(m, dm, ConcreteRNumber(3))
```
which then errors with
```
error: operand EnzymeAD/Reactant.jl#1 does not dominate this use
```
[mlir_dump.zip](https://github.com/user-attachments/files/27288369/mlir_dump.zip)
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.