# Is there anything like vmap to vectorize a computation

**URL:** <https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288>\
**Category:** GPU\
**Created:** [February 25, 2025, 12:32pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288 "2025-02-25T12:32:56Z")\
**Posts on this page:** 11\
**Page:** 1

<div class="post-metadata">

**Author:** ![Leander](https://avatars.discourse-cdn.com/v4/letter/l/bcef8e/32.png) [@Leander](https://discourse.julialang.org/u/Leander)\
**Post date:** [February 25, 2025, 12:32pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/1 "2025-02-25T12:32:56Z")

</div>

I want to evaluate using julia for a project and so far it was straightforward. The setting is the following, I have a bunch of functions (moment-generating functions over different parameters), for which I want to compute all the taylor coeffcients up to a certain order. My first try looks roughly like this:

```julia
using TaylorSeries
using LinearAlgebra
using CUDA

function moment_generating_function(u, vs)
    # more complicated in reality
    return simplex_volume(vs) * norm(u)
end

function simplex_volume(vs)
    """
    Calculate the volume of a simplex given its vertices.

    Parameters:
    vs (Matrix): A matrix of size (n+1, n) representing the vertices of the simplex,
                 where n is the dimension of the simplex.

    Returns:
    Float64: The volume of the simplex.
    """
    n = size(vs, 2)
    return abs(det(vs[2:end, :] .- vs[1, :])) / factorial(n)
end

function moment_generating_taylor_coefficients(vs, order)
    """
    Calculate the Taylor series expansion of the moment generating function of a simplex around 0.
    """
    d = size(vs, 2)
    u = set_variables("u1 u2", numvars=d, order=order)
    println("Taylor series variables: ", u)
    F = moment_generating_function(u, vs)
    c = F.coeffs
    println("Taylor series coefficient size" , size(c))
    return c
end

# Main script
df = [1.0 1.0; 2.0 5.0; 3.0 2.0]
println(size(df))
order = 7
taylor_coeffs = moment_generating_taylor_coefficients(df, order)

```

so far, this works great! But I have up to a few thousands of these simplices, and therefore I want to compute the taylor coefficients on the gpu. In torch, I would now call `vmap(lambda df: moment_generating_taylor_coefficients(df, order), axis=0)(simplices)` to map the `moment_generating_taylor_coefficients` along `axis=0`. How can I achieve something similar in julia? I am a bit lost with all the packages.

I guess I also need to change the data-type of `set_variables`.

I looked into gpu kernels but that looked extremely low-level.

tbh, I even struggle to compute the volume of the determinant of a batch of simplices. Afaik there’s no batch-dimension in the determinant function and as far as I know the for loops are not vectorized (otherwise I could just loop here).

---

<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:** [February 25, 2025, 2:06pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/2 "2025-02-25T14:06:19Z")

</div>

looking at TaylorSeries package, there is no external to CUDA or other gpu packages and nothing in docs about it, same for TaylorDiff, so, for now, I guess you can’t ? Note that you can make your code a lottt faster using

```julia
abs(det(@view(vs[2:end, :]) .- @view(vs[1, :]))) / factorial(n)

```

or even better, use a loop (if you stay on cpu) or AcceleratedKernels.jl can help on gpu but not for TaylorSeries.

About your code note that norm is not compatible with TaylorSeries (F is a float), you should use sqrt( sum(abs2,u) ) and do not diff around 0.

---

<div class="post-metadata">

**Author:** ![Leander](https://avatars.discourse-cdn.com/v4/letter/l/bcef8e/32.png) [@Leander](https://discourse.julialang.org/u/Leander)\
**Post date:** [February 25, 2025, 2:57pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/3 "2025-02-25T14:57:26Z")

</div>

thanks for the tip. Actually, judging from the issues, TaylorSeries is compatible? Looks like it broke recently though: [Fix from #44 to work with CUDA arrays Does Not Work in Latest Version · Issue #93 · JuliaDiff/TaylorDiff.jl · GitHub](https://github.com/JuliaDiff/TaylorDiff.jl/issues/93)

---

<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:** [February 25, 2025, 3:02pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/4 "2025-02-25T15:02:20Z")

</div>

It works fine on gpu just tested, I have a little issue with det indexing for some reason but if you replace det with something else it works

---

<div class="post-metadata">

**Author:** ![Leander](https://avatars.discourse-cdn.com/v4/letter/l/bcef8e/32.png) [@Leander](https://discourse.julialang.org/u/Leander)\
**Post date:** [February 25, 2025, 3:31pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/5 "2025-02-25T15:31:44Z")

</div>

How have you tried it? I tried it like in the issue and get the same error. Maybe I need to switch to another package like AcceleratedKernels.jl you mentioned

---

<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:** [February 25, 2025, 3:54pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/6 "2025-02-25T15:54:55Z")

</div>

this works

```julia
using TaylorSeries
using LinearAlgebra
using CUDA

function moment_generating_function(u, vs)
    # more complicated in reality
    return simplex_volume(vs) * sum(abs2,u)
end

function simplex_volume(vs)
    """
    Calculate the volume of a simplex given its vertices.

    Parameters:
    vs (Matrix): A matrix of size (n+1, n) representing the vertices of the simplex,
                 where n is the dimension of the simplex.

    Returns:
    Float64: The volume of the simplex.
    """
    n = size(vs, 2)
    return abs(sum(@views(vs[2:end, :] .- vs[1, :]))) / factorial(n)
end

function moment_generating_taylor_coefficients(vs, order)
    """
    Calculate the Taylor series expansion of the moment generating function of a simplex around 0.
    """
    d = size(vs, 2)
    u = set_variables("u1 u2", numvars=d, order=order)
    println("Taylor series variables: ", u)
    F = moment_generating_function(u, vs)
    # println("Taylor series coefficient size" , size(c))
    return F
end

# Main script
df = cu([1.0 1.0; 2.0 5.0; 3.0 2.0])
println(size(df))
order = 7
taylor_coeffs = moment_generating_taylor_coefficients(df, order)

```

but its not what you wanted to do since I removed det and norm which broke the code. This is closer but I don’t know about the performance (note that I diff around 1)

```julia
using TaylorSeries
using LinearAlgebra
using CUDA

function moment_generating_function(u, vs)
    # more complicated in reality
    return simplex_volume(vs) * sqrt(sum(abs2,u))
end

function simplex_volume(vs)
    """
    Calculate the volume of a simplex given its vertices.

    Parameters:
    vs (Matrix): A matrix of size (n+1, n) representing the vertices of the simplex,
                 where n is the dimension of the simplex.

    Returns:
    Float64: The volume of the simplex.
    """
    n = size(vs, 2)
    return abs(det(@views(vs[2:end, :] .- vs[1, :]))) / factorial(n)
end

function moment_generating_taylor_coefficients(vs, order)
    """
    Calculate the Taylor series expansion of the moment generating function of a simplex around 0.
    """
    d = size(vs, 2)
    u = set_variables("u1 u2", numvars=d, order=order)
    println("Taylor series variables: ", u)
    CUDA.@allowscalar F = moment_generating_function(1 .+u, vs)
    # println("Taylor series coefficient size" , size(c))
    return F
end

# Main script
df = cu([1.0 1.0; 2.0 5.0; 3.0 2.0])
println(size(df))
order = 7
CUDA.@time taylor_coeffs = moment_generating_taylor_coefficients(df, order)

```

---

<div class="post-metadata">

**Author:** ![Leander](https://avatars.discourse-cdn.com/v4/letter/l/bcef8e/32.png) [@Leander](https://discourse.julialang.org/u/Leander)\
**Post date:** [February 25, 2025, 4:07pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/7 "2025-02-25T16:07:59Z")

</div>

oh wow, that’s progress! Thank you! really cool.

I guess now I just have to figure out how to go from a single simplex (of shape (n+1,n)) to a batch of simplices (b, n+1, n).

---

<div class="post-metadata">

**Author:** ![danielwe](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/danielwe/32/35657_2.png) [@danielwe](https://discourse.julialang.org/u/danielwe)\
**Post date:** [February 25, 2025, 5:02pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/8 "2025-02-25T17:02:48Z")

</div>

GPU arrays implement generic julia functions as much as possible, including the broadcasting and mapreduce machinery, reducing the need for dedicated constructs like `vmap`. I haven’t tried it, but in this case, my first attempt would be the following:

```julia
mapslices(df -> moment_generating_taylor_coefficients(df, 7), simplices; dims=3)

```

where `simplices` has shape `(n + 1, n, b)`, with the batch dimension last such that the slices are contiguous in memory. (Julia arrays are column major, i.e., dimensions are ordered from “inner” to “outer”, opposite to python.)

---

<div class="post-metadata">

**Author:** ![Leander](https://avatars.discourse-cdn.com/v4/letter/l/bcef8e/32.png) [@Leander](https://discourse.julialang.org/u/Leander)\
**Post date:** [February 25, 2025, 5:53pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/9 "2025-02-25T17:53:31Z")

</div>

according to this:

> [@How to do mapslices() in parallel for 3D arrays](https://discourse.julialang.org/t/how-to-do-mapslices-in-parallel-for-3d-arrays/100484/4):
>
> Update my solutions here. Shorter but not efficient enough solution. This solution aligns each 3D array layer in 2D array, then computes the result of 2D array multiplication and last converts back to 3D array. using CUDA A = CUDA.rand(4, 4, 3) B = CUDA.rand(4, 4) C = reshape(permutedims(A, [1, 3, 2]), size(A, 1) \* size(A, 3), slight_smile \* B C = permutedims(reshape(C, size(A, 1), size(A, 3), :), [1, 3, 2]) Longer but more efficient solution. Use for loop to get each entry of the final 3D array. Thi…

this is not really working afaik

---

<div class="post-metadata">

**Author:** ![danielwe](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/danielwe/32/35657_2.png) [@danielwe](https://discourse.julialang.org/u/danielwe)\
**Post date:** [February 25, 2025, 6:04pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/10 "2025-02-25T18:04:31Z")

</div>

I see, that makes sense, mapping and broadcasting can’t replace, as @maleadt said in that thread,

> a tensor compiler that knows how to analyze the nested array ops and generate a single kernel

This sounds similar to what Reactant.jl aims to become. I haven’t tried it and don’t know whether it could solve your problem today.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [February 25, 2025, 6:48pm UTC](https://discourse.julialang.org/t/is-there-anything-like-vmap-to-vectorize-a-computation/126288/11 "2025-02-25T18:48:42Z")

</div>

If your `n` is fairly small (\<10 maybe) then you can try a CuVector of SMatrix, using StaticArrays.jl . That’s the same memory as a size `(n + 1, n, b)` CuArray, but these static arrays are like scalars (or complex numbers) as far as CUDA.jl is concerned, which allows some `mapslices`-like things to be done.