# Julia's Broadcast vs Jax's vmap

**URL:** https://discourse.julialang.org/t/julias-broadcast-vs-jaxs-vmap/38990
**Category:** Internals & Design
**Created:** [May 7, 2020, 10:12pm UTC](https://discourse.julialang.org/t/julias-broadcast-vs-jaxs-vmap/38990 "2020-05-07T22:12:37Z")
**Posts on this page:** 1
**Showing post:** 9

<div class="post-metadata">

### Author: ![darsnack](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/darsnack/32/10144_2.png) [@darsnack](https://discourse.julialang.org/u/darsnack)
#### Post date: [May 7, 2020, 11:14pm UTC](https://discourse.julialang.org/t/julias-broadcast-vs-jaxs-vmap/38990/9 "2020-05-07T23:14:54Z")

</div>

I ran the benchmarks on the GPU:

```julia
using BenchmarkTools, CuArrays
using LinearAlgebra: dot

D = 10^3
BS = 10^2

x = randn(D)
X = randn(D, BS)
y = randn(D)
cX = cu(X)
cy = cu(y)
Xt = permutedims(X)
cXt = cu(Xt)

dot(x, y)
dot(cu(x), cy)

broadcast_dot(X, y) = [dot(x, y) for x in eachslice(X; dims = 2)]
matmul_dot(Xt, y) = Xt * y

```

Now running on the CPU:

```julia
@btime broadcast_dot($X, $y)
16.652 μs (108 allocations: 6.56 KiB)

@btime matmul_dot($Xt, $y)
13.867 μs (1 allocation: 896 bytes)

```

And on the GPU:

```julia
@btime CuArrays.@sync broadcast_dot($cX, $cy)
321.091 ms (208 allocations: 8.45 KiB)

@btime CuArrays.@sync matmul_dot($cXt, $cy)
238.609 μs (8 allocations: 208 bytes)

```

Of note is that the following definition did not work:

```julia
broadcast_dot(X, y) = dot.(eachslice(X; dims = 2), Ref(y))

```

---

_[View the full topic](https://discourse.julialang.org/t/julias-broadcast-vs-jaxs-vmap/38990)._