EnzymeAD / EnzymeAD/Enzyme-JAX

autodiff through while loops with state-dependent break

Open
#2,565 0 comments 0 reactions 1 assignee Claimed by @Pangoraw View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.