Hi all,

I’m fairly new to Julia, coming from Python + JAX. As such, I’ve become accustomed to writing vectorized expressions to avoid expensive `for`

loops – however, I understand that this isn’t something that should be a concern in Julia. I do like functional programming in general though, so I have really enjoyed JAX’s `vmap`

functionality for expressing such computations. I have been trying to write my Julia code in a similar fashion, although occasionally I get stuck and it feels like I’m fighting with Julia just to write this way, which of course I want to avoid.

This afternoon I found what I think is a nice solution to a problem I was working on. I have a batch of `N`

vectors of dimension `A`

, and it’s represented by a `Matrix m`

of size `(A, N)`

. For each vector, I want to compute the softmax and sample a categorical variable parameterized by the softmax as the category probabilities. A first attempt was something like

```
using Distributions
using Random
[rand(Categorical(p)) for p in eachcol(m)]
```

but this doesn’t work – I get an exception like

```
MethodError: Cannot `convert` an object of type Vector{Float64} to an object of type SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}
```

Eventually I realized I could fix this with

```
[rand(Categorical(collect(p))) for p in eachcol(m)]
```

but this is fairly hard to read. Finally I settled on the following,

```
map(rand ∘ Categorical ∘ collect, eachcol(m))
```

which I find much easier to read. However, I rarely see code like this written in Julia, and I’m concerned that there are consequences of this code that I’m unaware of.

How would you write write this code, and is there anything outright bad with the way that I wrote it?