# Lux.jl vs Jax
**URL:** https://discourse.julialang.org/t/lux-jl-vs-jax/119654
**Category:** General Usage
**Tags:** machine-learning, jax, lux
**Created:** [September 20, 2024, 6:22pm UTC](https://discourse.julialang.org/t/lux-jl-vs-jax/119654 "2024-09-20T18:22:57Z")
**Posts on this page:** 7
**Page:** 1
### Author:  [@ijk](https://discourse.julialang.org/u/ijk)
#### Post date: [September 20, 2024, 6:22pm UTC](https://discourse.julialang.org/t/lux-jl-vs-jax/119654/1 "2024-09-20T18:22:57Z")
I’ve started using Julia this semester, and would like to reimplement some ML models from Jax to Lux.jl. I have three concerns:
1. I noticed there aren’t constructs like `vmap`. My existing Jax code has quite a few vmap calls. It’s crucial to take advantage of the GPU.
2. How fast is Lux.jl compared to Jax implementation for moderately beefy models?
3. I have a hard time designing structs which are to be vectorized. For example, I could have
```julia
struct Foo
x::Float32
end
xs = Foo[...]
```
or
```julia
struct Foo
x::Vector{Float32}
end
```
I prefer the first one since it is easier to reason with, but I believe the second is necessary for good GPU performance. Do I need to make this trade-off or am I missing something?
1. GPU Samplers. The Distributions.jl package seems to implement distributions on the CPU, but I don’t see anything about GPU sampling.
Thank you!
---
### Author:  [@marius311](https://discourse.julialang.org/u/marius311)
#### Post date: [September 20, 2024, 6:43pm UTC](https://discourse.julialang.org/t/lux-jl-vs-jax/119654/2 "2024-09-20T18:43:59Z")
Welcome! Answers to 2 of your questions:
> 1. I noticed there aren’t constructs like `vmap`. My existing Jax code has quite a few vmap calls. It’s crucial to take advantage of the GPU.
No `vmap` in Julia currently unfortunately. Just try putting your loop on the outer most level with a normal `map` rather than pushing to the inner-most loop like `vmap`. As long as your inner-most loops already saturate the GPU, it won’t really matter (if they don’t though you do have to think more than the easy solution `vmap` offers)
> 1. I have a hard time designing structs which are to be vectorized. For example, I could have
A unique thing Julia can do that Jax can’t is put arrays of custom structs on GPU as long as the the structs are concrete (all types known at compile time). So e.g. this is totally valid and performant
```julia
julia> struct Foo
x::Float32
end
julia> cu(Foo.([1,2,3]))
3-element CuArray{Foo, 1, CUDA.Mem.DeviceBuffer}:
Foo(1.0f0)
Foo(2.0f0)
Foo(3.0f0)
```
so if you find code is easier to write/reason about that way, go for it.
---
### Author:  [@ijk](https://discourse.julialang.org/u/ijk)
#### Post date: [September 20, 2024, 6:59pm UTC](https://discourse.julialang.org/t/lux-jl-vs-jax/119654/3 "2024-09-20T18:59:52Z")
Thank you! I didn’t know you can do that with structs. Could you elaborate on Q1 though, maybe with a simple example?
---
### Author:  [@marcobonici](https://discourse.julialang.org/u/marcobonici)
#### Post date: [September 20, 2024, 7:25pm UTC](https://discourse.julialang.org/t/lux-jl-vs-jax/119654/4 "2024-09-20T19:25:20Z")
On Q1, what do you want to do? Give the input NN a batch of input arrays simoultaneously?
---
### Author:  [@jar1](https://discourse.julialang.org/u/jar1)
#### Post date: [September 20, 2024, 7:29pm UTC](https://discourse.julialang.org/t/lux-jl-vs-jax/119654/5 "2024-09-20T19:29:39Z")
> [@ijk](#):
>
> I have a hard time designing structs which are to be vectorized.
I like StructArrays.jl which does the “struct of arrays” arrangement.
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [September 20, 2024, 8:38pm UTC](https://discourse.julialang.org/t/lux-jl-vs-jax/119654/6 "2024-09-20T20:38:59Z")
For Q1, I think the closest thing we’ll have in Julia is Reactant.jl, but it’s still a very very early prototype. @avikpal and @wsmoses would be able to tell us more
---
### Author:  [@avikpal](https://discourse.julialang.org/u/avikpal)
#### Post date: [September 21, 2024, 6:12pm UTC](https://discourse.julialang.org/t/lux-jl-vs-jax/119654/7 "2024-09-21T18:12:02Z")
> [@ijk](#):
>
> How fast is [Lux.jl](https://juliahub.com/ui/Packages/General/Lux) compared to Jax implementation for moderately beefy models?
Generally beefy models on CUDA are just cuDNN calls internally, so it should be comparable performance. It is a bit hard to say for sure before looking at the model, but if there is a noticeable slowdown, open an issue or make a post here or in github and we can take a look.
For CPU it would be slower if you have conv routines (smaller models are actually faster than pytorch but that is not what you are interested in), and for ROCm again it depends because we don’t have all the bindings. Metal and oneAPI are experimental and performance there is quite bad atm.
But overall, the eventual goal is to be able to compile via Reactant, and if you take a look at the Reactant and Lux repo, we are actively working on an easy way to take a Lux model and make it faster.