# Gradient through categorical distribution slow and has many allocations

**URL:** <https://discourse.julialang.org/t/gradient-through-categorical-distribution-slow-and-has-many-allocations/43690>\
**Category:** Performance\
**Created:** [July 26, 2020, 1:21am UTC](https://discourse.julialang.org/t/gradient-through-categorical-distribution-slow-and-has-many-allocations/43690 "2020-07-26T01:21:18Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![scottjordan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/scottjordan/32/26096_2.png) [@scottjordan](https://discourse.julialang.org/u/scottjordan)\
**Post date:** [July 26, 2020, 1:21am UTC](https://discourse.julialang.org/t/gradient-through-categorical-distribution-slow-and-has-many-allocations/43690/1 "2020-07-26T01:21:18Z")

</div>

I am trying have performant gradients computations for a function the defines a conditional distribution over a discrete space. More specifically, I want the gradient of the following function:

```julia
import Flux
import Zygote
using Distributions
using BenchmarkTools

function logprob(W, x, a)
    probs = Flux.softmax(W' * x))
    d = Categorical(probs)
    return logpdf(d,a)
end

x = collect(Float64, 1:100)
w = zeros(Float64, (100,2))
@btime (w->Zygote.gradient(logprob(w,x,1),w)
# running stats: 150.143 μs (608 allocations: 29.55 KiB)

```

By removing the categorical distribution I tend to see about two orders of magnitude decrease in computation time.

```julia
function logprob2(W, x, a)
    probs = Flux.softmax(W' * x))
    return log(probs[a])
end
@btime (w->Zygote.gradient(logprob2(w,x,1),w)
# running stats: 1.031 μs (26 allocations: 4.53 KiB)

```

My questions are: what is causing the slow down and is it possible to speed up the process while still leveraging Distributions.jl functions?