# Trying to implement VAE using Lux and reactant

**URL:** <https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353>\
**Category:** General Usage\
**Tags:** reactant\
**Created:** [January 2, 2025, 10:06am UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353 "2025-01-02T10:06:40Z")\
**Posts on this page:** 13\
**Page:** 1

<div class="post-metadata">

**Author:** ![yolhan\_mannes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yolhan_mannes/32/220485_2.png) [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)\
**Post date:** [January 2, 2025, 10:06am UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/1 "2025-01-02T10:06:40Z")

</div>

Just had a lot of fun implementing VAE (Variational AutoEncoder) with Lux and Reactant, it barely works, little summary of errors,

Reactant : errors with:

```julia
'stablehlo.transpose' op using value defined outside the region  
ERROR: "failed to run pass manager on module"  

```

CPU with Enzyme errors with

```julia
Duplicated(Decoder,RefValue) error  

```

note : gradient can be calculated when using Reactant, but not when using cpu, however training still fail

works with AutoZygote() on cpu, did not try on gpu.

repo : [GitHub - yolhan83/MLX\_exemples\_julia\_reactant: trying to implement some of the "mlx-examples" repo code](https://github.com/yolhan83/MLX_exemples_julia_reactant)

I know I should make a mwe from that but that may be hard, will see. If anyone has ideas they are welcome.

---

<div class="post-metadata">

**Author:** ![simeonschaub](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/simeonschaub/32/216566_2.png) [@simeonschaub](https://discourse.julialang.org/u/simeonschaub)\
**Post date:** [January 2, 2025, 10:43am UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/2 "2025-01-02T10:43:09Z")

</div>

> [@yolhan\_mannes](#):
>
> Reactant : errors with:
> 
> ```julia-auto
> 'stablehlo.transpose' op using value defined outside the region  
> ERROR: "failed to run pass manager on module"  
> 
> ```

I encountered similar errors before, I believe it typically means you are mixing different floating point precisions, which Reactant cannot handle ATM. Looking at your code, it looks like [this line](https://github.com/yolhan83/MLX_exemples_julia_reactant/blob/5bd9ae5b428054c74e83a65f20be0183fb7e16df/cvae/main.jl#L33) specifically could be the issue, as it’s using `0.5` which is of type `Float64` and it looks like the rest of your code is using `Float32`. Otherwise it would help to see the whole stacktrace

---

<div class="post-metadata">

**Author:** ![yolhan\_mannes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yolhan_mannes/32/220485_2.png) [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)\
**Post date:** [January 2, 2025, 11:46am UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/3 "2025-01-02T11:46:45Z")

</div>

thanks, yes that could be bad, still not working same error, here is the full error

[err.jl](https://discourse.julialang.org/uploads/short-url/58sllpS1Zite9tKGtyitiaUakI9.jl) (85.0 KB)

btw loss and loss gradient compile fine its really in the optimisation process

---

<div class="post-metadata">

**Author:** ![simeonschaub](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/simeonschaub/32/216566_2.png) [@simeonschaub](https://discourse.julialang.org/u/simeonschaub)\
**Post date:** [January 2, 2025, 12:00pm UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/4 "2025-01-02T12:00:29Z")

</div>

If you look at the stacktrace, you can still see `ConcreteRNumber{Float64}` in there, meaning you are still promoting to `Float64` somewhere. Since you say this is only in the optimization process, my best guess is that you need to make your learning rate in the Adam optimizer a `Float32` as well, so replace `1e-3` with `1f-3`

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [January 2, 2025, 2:31pm UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/5 "2025-01-02T14:31:08Z")

</div>

The transpose issue is definitely a bug and shouldn’t happen (regardless of mixed precision or not, which should be fine?)

please open an issue with a reproducer (ideally with any amount of reduction for where it’s coming from)

Also cc @avikpal

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [January 2, 2025, 2:49pm UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/6 "2025-01-02T14:49:05Z")

</div>

That won’t help Optimisers.jl currently stores everything in Float64 (see [Type Constraints in the Rule Structs · Issue #205 · FluxML/Optimisers.jl · GitHub](https://github.com/FluxML/Optimisers.jl/issues/205)).

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [January 2, 2025, 2:55pm UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/7 "2025-01-02T14:55:19Z")

</div>

Is this on the latest releases of Lux and Reactant?

I also happen to have a partial implementation of CVAE from the MLX repo [Lux.jl/examples/ConditionalVAE/main.jl at ap/cvae2 · LuxDL/Lux.jl · GitHub](https://github.com/LuxDL/Lux.jl/blob/ap/cvae2/examples/ConditionalVAE/main.jl), but probably needs to be updated

---

<div class="post-metadata">

**Author:** ![yolhan\_mannes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yolhan_mannes/32/220485_2.png) [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)\
**Post date:** [January 2, 2025, 4:33pm UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/8 "2025-01-02T16:33:54Z")

</div>

Yes it is the latest on both and your code looks very similar to mine, is it working fine ?

---

<div class="post-metadata">

**Author:** ![yolhan\_mannes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yolhan_mannes/32/220485_2.png) [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)\
**Post date:** [January 2, 2025, 4:35pm UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/9 "2025-01-02T16:35:39Z")

</div>

Thank you I was getting crazy trying to find where those are, I will try to make the mwe

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [January 2, 2025, 4:49pm UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/10 "2025-01-02T16:49:00Z")

</div>

yes, it’s training fine for the most part. There are some of the usual issues of VAEs with NaNs, which I am trying to sort out

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [January 3, 2025, 5:59pm UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/11 "2025-01-03T17:59:01Z")

</div>

The final functional version has been merged [Lux.jl/examples/ConditionalVAE/main.jl at main · LuxDL/Lux.jl · GitHub](https://github.com/LuxDL/Lux.jl/blob/main/examples/ConditionalVAE/main.jl)

 ![conditional_vae](https://global.discourse-cdn.com/julialang/original/3X/5/f/5fe66b9f44cd5a9f948e4969766f9fc8f5ef651c.png)

---

<div class="post-metadata">

**Author:** ![yolhan\_mannes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yolhan_mannes/32/220485_2.png) [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)\
**Post date:** [January 3, 2025, 6:50pm UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/12 "2025-01-03T18:50:41Z")

</div>

Thank you I’ve made mine work too by putting layers together in Chain instead of having them separate in the struct did not try to make the mwe yet though will see

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [January 8, 2025, 9:02pm UTC](https://discourse.julialang.org/t/trying-to-implement-vae-using-lux-and-reactant/124353/13 "2025-01-08T21:02:21Z")

</div>

the transpose issue has been resolved in the latest releases of Reactant (v0.2.17) and Lux (v1.5)