# Use Enzyme in flux
**URL:** https://discourse.julialang.org/t/use-enzyme-in-flux/98352
**Category:** Machine Learning
**Tags:** flux
**Created:** [May 4, 2023, 10:34pm UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352 "2023-05-04T22:34:30Z")
**Posts on this page:** 11
**Page:** 1
### Author:  [@andferrari](https://discourse.julialang.org/u/andferrari)
#### Post date: [May 4, 2023, 10:34pm UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/1 "2023-05-04T22:34:30Z")
Hi, I have a complex loss with mutating arrays unsupported by Zygote.
Is it possible to use Enzyme.jl with Flux.jl
---
### Author:  [@e3c6](https://discourse.julialang.org/u/e3c6)
#### Post date: [May 4, 2023, 10:36pm UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/2 "2023-05-04T22:36:26Z")
I’m also interested in the answer. Would like to see an example of Enzyme usage with Flux / Lux.
---
### Author:  [@cortner](https://discourse.julialang.org/u/cortner)
#### Post date: [May 5, 2023, 12:05am UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/3 "2023-05-05T00:05:00Z")
You can write a custom rrule and use Enzyme to implement that
---
### Author:  [@andferrari](https://discourse.julialang.org/u/andferrari)
#### Post date: [May 5, 2023, 11:01am UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/4 "2023-05-05T11:01:38Z")
thank you for your answer.
Could you please provide a custom rule basic example?
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [May 5, 2023, 11:37am UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/5 "2023-05-05T11:37:55Z")
There is one in the docs for the latest version: [Custom rules · Enzyme.jl](https://enzyme.mit.edu/julia/stable/generated/custom_rule/)
---
### Author:  [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [May 5, 2023, 12:31pm UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/6 "2023-05-05T12:31:46Z")
> [@e3c6](#):
>
> Would like to see an example of Enzyme usage with Flux / Lux.
```julia
using Enzyme
x = [2.0, 2.0]
bx = [0.0, 0.0]
y = [0.0,0.0]
using ComponentArrays, Lux, Random
rng = Random.default_rng()
Random.seed!(rng,100)
dudt2 = Lux.Chain(x -> x.^3,
Lux.Dense(2, 50, tanh),
Lux.Dense(50, 2))
p, st = Lux.setup(rng, dudt2)
function f(x::Array{Float64}, y::Array{Float64})
y .= dudt2(x, p, st)[1]
return nothing
end
Enzyme.autodiff(Reverse, f, Duplicated(x, bx), Duplicated(y, ones(2)))
function f2(x::Array{Float64})
dudt2(x, p, st)[1]
end
using Zygote
bx2 = Zygote.pullback(f2, x)[2](ones(2))[1]
bx
@show bx - bx2
#=
2-element Vector{Float64}:
-9.992007221626409e-16
-1.7763568394002505e-15
=#
```
on main, unreleased.
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [May 6, 2023, 3:03am UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/7 "2023-05-06T03:03:44Z")
I would caution this coding style however and strongly recommend passing in dudt p and st explicitly as (const) parameters to the autodiff rather than type unstably capturing them.
It is significant for performance among other things.
---
### Author:  [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [May 6, 2023, 3:13am UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/8 "2023-05-06T03:13:18Z")
Agreed. Just part of the demo that globals do work (dudt is just a function though)
---
### Author:  [@andferrari](https://discourse.julialang.org/u/andferrari)
#### Post date: [May 8, 2023, 7:51am UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/9 "2023-05-08T07:51:10Z")
Thank you for your help!
I understand that computation of the gradient vector is not a problem. But what I miss is knowing how to convert the gradient vector in order to use it in update!() for parameters update.
---
### Author:  [@Maysam\_Gholampour](https://discourse.julialang.org/u/Maysam_Gholampour)
#### Post date: [June 21, 2024, 1:06am UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/10 "2024-06-21T01:06:23Z")
this doesn’t work on GPU. Did I do sth wrong?
```julia
using Enzyme
using Lux, Random, LuxCUDA
rng = Random.default_rng()
Random.seed!(rng,100)
dudt2 = Lux.Chain(x -> x.^3,
Lux.Dense(2, 50, tanh),
Lux.Dense(50, 2))
gpu_dev = gpu_device()
p, st = Lux.setup(rng, dudt2) .|> gpu_dev
function f(x::T, y::T) where T
y .= dudt2(x, p, st)[1]
return nothing
end
x = [2.0f0, 2.0f0] |> gpu_dev
bx = [0.0f0, 0.0f0] |> gpu_dev
y = [0.0f0,0.0f0] |> gpu_dev
ones32 = ones(Float32, 2) |> gpu_dev
Enzyme.autodiff(Reverse, f, Duplicated(x, bx), Duplicated(y, ones32))
```
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [June 21, 2024, 9:17pm UTC](https://discourse.julialang.org/t/use-enzyme-in-flux/98352/11 "2024-06-21T21:17:57Z")
File an issue?
though also forewarning, Differentiating host-side code when accesses device memory (e.g. `sum(CuArray)` ) is not yet supported, but in progress. (see [FAQ · Enzyme.jl](https://enzyme.mit.edu/julia/stable/faq/#CUDA-support))