Skip to content

Enzyme.jl support on Julia 1.12+ #1323

@ChrisRackauckas-Claude

Description

@ChrisRackauckas-Claude

Summary

Enzyme.jl does not currently support Julia 1.12+. This issue tracks the workarounds implemented and items to revisit once Enzyme v1.12 support lands.

Upstream Issue

See: EnzymeAD/Enzyme.jl#2699

Current Status

Enzyme.jl displays the following warning on Julia 1.12:

Warning: Enzyme.jl support for Julia 1.12 is presently in progress.
         For the time being we recommend using 1.11 or LTS (1.10).

Affected Areas with MWEs

All MWEs below fail on Julia 1.12+ but work on Julia 1.11 and earlier.

1. automatic_sensealg_choice - EnzymeVJP not selected (Julia 1.12+)

The automatic_sensealg_choice function tries Enzyme first for in-place VJP computation. On Julia 1.12+, this fails and falls back to ReverseDiffVJP.

MWE (fails on Julia 1.12+, works on Julia 1.11):

using Enzyme, OrdinaryDiffEq

function f!(du, u, p, t)
    du[1] = p[1] * u[1] - p[2] * u[1] * u[2]
    du[2] = -p[3] * u[2] + p[4] * u[1] * u[2]
    return nothing
end

u0 = [1.0, 1.0]
p = [1.5, 1.0, 3.0, 1.0]
du = zero(u0)
t0 = 0.0

# This fails on Julia 1.12+, works on Julia 1.11
Enzyme.autodiff(
    Enzyme.Reverse, f!,
    Enzyme.Duplicated(du, copy(u0)),
    Enzyme.Duplicated(copy(u0), zero(u0)),
    Enzyme.Duplicated(copy(p), zero(p)),
    Enzyme.Const(t0)
)

Test: test/automatic_sensealg_choice.jl now marks EnzymeVJP as @test_broken on Julia 1.12+.

2. NonlinearProblem gradient computation (Julia 1.12+)

The NonlinearProblem tests in test/steady_state.jl use Enzyme for gradient computation.

MWE (fails on Julia 1.12+, works on Julia 1.11):

using Enzyme, NonlinearSolve, SciMLSensitivity

u0 = [0.0]
p = [2.0, 1.0]
prob = NonlinearProblem((du, u, p) -> du[1] = u[1] - p[1] + p[2], u0, p)

function test_loss(p, prob, alg)
    _prob = remake(prob, p = p)
    sol = sum(solve(_prob, alg, sensealg = SteadyStateAdjoint(autojacvec = ReverseDiffVJP())))
    return sol
end

# This fails on Julia 1.12+, works on Julia 1.11
dp = Enzyme.make_zero(p)
dprob = Enzyme.make_zero(prob)
Enzyme.autodiff(
    Enzyme.Reverse, test_loss, Enzyme.Active,
    Enzyme.Duplicated(p, dp),
    Enzyme.Duplicated(prob, dprob),
    Enzyme.Const(NewtonRaphson())
)

Test: test/steady_state.jl NonlinearProblem Enzyme tests are now conditionally skipped on Julia 1.12+.

3. Lux Neural Network ODE (Julia 1.12+)

The automatic sensealg choice test uses a Lux neural network ODE which also triggers Enzyme failures.

MWE (fails on Julia 1.12+, works on Julia 1.11):

using Lux, ComponentArrays, OrdinaryDiffEq, Enzyme, Random

rng = Random.default_rng()
ann = Lux.Chain(Lux.Dense(1, 32, tanh), Lux.Dense(32, 32, tanh), Lux.Dense(32, 1))
ps, st = Lux.setup(rng, ann)
p = ComponentArray(ps)
θ, ax = getdata(p), getaxes(p)

function dxdt!(dx, x, p, t)
    ps = ComponentArray(p, ax)
    dx[1] = x[2]
    dx[2] = first(ann([t], ps, st))[1]^3
    return nothing
end

x0 = [-4.0f0, 0.0f0]
du = zero(x0)
t0 = 0.0f0

# This fails on Julia 1.12+, works on Julia 1.11
Enzyme.autodiff(
    Enzyme.Reverse, dxdt!,
    Enzyme.Duplicated(du, copy(x0)),
    Enzyme.Duplicated(copy(x0), zero(x0)),
    Enzyme.Duplicated(copy(θ), zero(θ)),
    Enzyme.Const(t0)
)

TODO when Enzyme 1.12 support lands

  • Remove @test_broken from automatic_sensealg_choice.jl
  • Re-enable Enzyme tests in steady_state.jl NonlinearProblem section
  • Consider switching back to EnzymeVJP as default autojacvec on Julia 1.12+
  • Review if any other tests need to use Enzyme instead of fallback backends

Related Changes

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions