Changed gradient of LSE (#885)

TuringLang · Aug 9, 2019 · abbc8c0 · abbc8c0
1 parent 9f6f127
commit abbc8c0
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/core/ad.jl b/src/core/ad.jl
@@ -161,10 +161,9 @@ end
 import StatsFuns: logsumexp
 logsumexp(x::Tracker.TrackedArray) = Tracker.track(logsumexp, x)
 Tracker.@grad function logsumexp(x::Tracker.TrackedArray)
-    lse = logsumexp(Tracker.data(x)) 
-    se = exp(lse)
+    lse = logsumexp(Tracker.data(x))
     return lse,
-          Δ->(Δ .* exp.(x) ./ se,)
+          Δ->(Δ .* exp.(x .- lse),)
 end
 
 import StatsFuns: binomlogpdf