Added fix to LogSoftMax to handle large values without going to Inf

abeschneider · Jan 8, 2017 · bca9ec3 · bca9ec3
1 parent 0d6618e
commit bca9ec3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/Sources/Tensor/tensorops.swift b/Sources/Tensor/tensorops.swift
@@ -968,8 +968,10 @@ public func logsoftmax<S:Storage>
     // log[ exp(input) / sum(exp(input)) ]
     // = log(exp(input)) - log(sum(exp(input)))
     // = input - log(sum(exp(input)))
-    let s = Tensor<S>([sum(exp(input))])
-    let logsum = log(s)
+    // maxInput scales all values so we don't get Inf due to the exp
+    let maxInput = Tensor<S>([max(input)])
+    let s = Tensor<S>([sum(exp(input - maxInput))])
+    let logsum = maxInput + log(s)
     sub(input, logsum, result: result)
 }