# Softmax using Tullio

**URL:** https://discourse.julialang.org/t/softmax-using-tullio/89838
**Category:** General Usage
**Created:** [November 6, 2022, 7:17am UTC](https://discourse.julialang.org/t/softmax-using-tullio/89838 "2022-11-06T07:17:30Z")
**Posts on this page:** 3
**Page:** 1

<div class="post-metadata">

### Author: ![minetest2048](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/minetest2048/32/45961_2.png) [@minetest2048](https://discourse.julialang.org/u/minetest2048)
#### Post date: [November 6, 2022, 7:17am UTC](https://discourse.julialang.org/t/softmax-using-tullio/89838/1 "2022-11-06T07:17:30Z")

</div>

I want to port my DIY softmax function

```julia
softmax(x) = exp.(x .- maximum(x)) / sum(exp.(x .- maximum(x)))

```

to Tullio einsum so that it runs faster and gives fast symbolic gradients

I split it into multiple lines, and it works:

```julia
function softmax_einsum(x)
    maxx = maximum(x)
    @tullio sumx := exp(x[i] - maxx) verbose=false
    @tullio ret[i] := exp(x[i] - maxx) / sumx verbose=false
end

```

It works

```julia
julia> t1 = [1.0,1,1,1]
4-element Vector{Float64}:
 1.0
 1.0
 1.0
 1.0

julia> softmax(t1)
4-element Vector{Float64}:
 0.25
 0.25
 0.25
 0.25

julia> NNlib.softmax(t1)
4-element Vector{Float64}:
 0.25
 0.25
 0.25
 0.25

julia> softmax_einsum(t1)
4-element Vector{Float64}:
 0.25
 0.25
 0.25
 0.25

```

But it gives different jacobian

```julia
julia> t1 = [1.0,1,1,1]
4-element Vector{Float64}:
 1.0
 1.0
 1.0
 1.0

julia> Zygote.jacobian(NNlib.softmax, t1)
([0.1875 -0.0625 -0.0625 -0.0625; -0.0625 0.1875 -0.0625 -0.0625; -0.0625 -0.0625 0.1875 -0.0625; -0.0625 -0.0625 -0.0625 0.1875],)

julia> Zygote.jacobian(softmax, t1)
([0.1875 -0.0625 -0.0625 -0.0625; -0.0625 0.1875 -0.0625 -0.0625; -0.0625 -0.0625 0.1875 -0.0625; -0.0625 -0.0625 -0.0625 0.1875],)

julia> Zygote.jacobian(softmax_einsum, t1)
([0.25 0.0 0.0 0.0; 0.0 0.25 0.0 0.0; 0.0 0.0 0.25 0.0; 0.0 0.0 0.0 0.25],)

```

What happened?

---

<div class="post-metadata">

### Author: ![rveltz](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rveltz/32/2707_2.png) [@rveltz](https://discourse.julialang.org/u/rveltz)
#### Post date: [November 6, 2022, 8:11am UTC](https://discourse.julialang.org/t/softmax-using-tullio/89838/2 "2022-11-06T08:11:49Z")

</div>

you are computing `exp(...)` twice.

---

<div class="post-metadata">

### Author: ![albheim](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albheim/32/34660_2.png) [@albheim](https://discourse.julialang.org/u/albheim)
#### Post date: [November 6, 2022, 9:12am UTC](https://discourse.julialang.org/t/softmax-using-tullio/89838/3 "2022-11-06T09:12:28Z")

</div>

Yeah, im not sure why you get the wrong thing with tullio, but you can improve your implementation before that.

First I dont think the maximum does much if you are looking at it symbolically, it might have some numerical purpose but the answer should be the same without it.  
And the as was mentioned above you calculate exp of all elements twice which is a large waste.