Forcing Lux to use Enzyme instead of Zygote?

Specific Domains Machine Learning
1

Hi everyone.

I’m confused.
The example that used to work @GitHub - SciML/NeuralPDE.jl: Physics-Informed Neural Networks (PINN) and Deep BSDE Solvers of Differential Equations for Scientific Machine Learning (SciML) accelerated simulation
does not work anymore because of the common error of Zygote

ERROR: Mutating arrays is not supported – called setindex!(Vector{Float64}, …)
This error occurs when you ask Zygote to perform operations that mutate the elements of
the elements of arrays in place (e.g. setting values with x .= …).

Possible fixes:

  • avoid mutating operations (preferred)
  • or read the documentation and solutions for this bug
    Limitations · Zygote

I tried to see how to fix this but didn’t succeed.

So I thought, what about forcing Lux to use Enzyme instead of Zygote?
So I added the code example using Enzyme.

but nothing happened at all.

If any of you can help me, ask Lux to use Enzyme instead of Zygote. Maybe I can see why this is happening when this example used to run perfectly.

Naceur

2

What example?

3

Hi Chris :slight_smile:

The one on the front page of NeuralPDE : This one


using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers
import ModelingToolkit: Interval, infimum, supremum

@parameters x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2

# 2D PDE
eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y)

# Boundary conditions
bcs = [u(0, y) ~ 0.0, u(1, y) ~ 0,
    u(x, 0) ~ 0.0, u(x, 1) ~ 0]
# Space and time domains
domains = [x ∈ Interval(0.0, 1.0),
    y ∈ Interval(0.0, 1.0)]
# Discretization
dx = 0.1

# Neural network
dim = 2 # number of dimensions
chain = Lux.Chain(Dense(dim, 16, Lux.σ), Dense(16, 16, Lux.σ), Dense(16, 1))

discretization = PhysicsInformedNN(chain, QuadratureTraining())

@named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
prob = discretize(pde_system, discretization)

callback = function (p, l)
    println("Current loss is: $l")
    return false
end

res = Optimization.solve(prob, ADAM(0.1); callback = callback, maxiters = 4000)
prob = remake(prob, u0 = res.minimizer)
res = Optimization.solve(prob, ADAM(0.01); callback = callback, maxiters = 2000)
phi = discretization.phi

xs, ys = [infimum(d.domain):(dx / 10):supremum(d.domain) for d in domains]
analytic_sol_func(x, y) = (sin(pi * x) * sin(pi * y)) / (2pi^2)

u_predict = reshape([first(phi([x, y], res.minimizer)) for x in xs for y in ys],
                    (length(xs), length(ys)))
u_real = reshape([analytic_sol_func(x, y) for x in xs for y in ys],
                 (length(xs), length(ys)))
diff_u = abs.(u_predict .- u_real)

using Plots
p1 = plot(xs, ys, u_real, linetype = :contourf, title = "analytic");
p2 = plot(xs, ys, u_predict, linetype = :contourf, title = "predict");
p3 = plot(xs, ys, diff_u, linetype = :contourf, title = "error");
plot(p1, p2, p3)
4

Chris

As soon as I will go to my office I will downgrade Zygote and see what will happen
I already downgraded Lux NeuralPDE and all the other packages but nothing happened all the codes and not only this example do not run

I think it has to do with the late upgrade of Zygote

5

See:

6

I think the new ChainRules was released which fixes this. Try upgrading?

7

Chris

Yeeeeeeeeeeeeeeeeeeeeees it works now :slight_smile: :star_struck: :beers:

Many thanks Chris !

Because I really started to lose confidence in my self to still have the ability to learn

:slight_smile:

All what I know comes from you and your team Chris

Thank you a lot

Naceur

8

Thank @Mason for the bug squashing.

9

Thanks @Mason :slight_smile:

I’m testing now my own small code
Right now ! It seems working but instable : the code start much better than before, the gradient descent is heading towards the solution but instable yet !
It could come from my inputs or grid sampling ; I am using NeuralPDE for heat equation, but the initial condition is somehow hard ! I will tell you guys if this come from the difficulty of my setup or there is something unusual

10

You can also blame me for the bug creation, so both “sorry” and “your welcome” :sweat_smile:

11

What Chris ! :slight_smile:

We can only thank Chris and all those who help demystify so many wonderful things.
Which will undoubtedly contribute to a better world.

Thank you guys :slight_smile: