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:

2 Likes
6

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

1 Like
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

1 Like
8

Thank @Mason for the bug squashing.

2 Likes
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

1 Like
10

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

1 Like
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: