Numerically stable logit pred

[dr-david]

The get_logit_predictions function was computing $\text{logit}(\text{softmax}(x))$ which is not stable for values of $\text{softmax}(x) \approx 0 \text{ or } \approx 1$

I rewrote it as:

$$ \text{logit}(\text{softmax}(x)_i) = \log(\text{softmax}(x)_i) - \log(1 - \text{softmax}(x)_i) $$

Expanding further:

$$ \text{logit}(\text{softmax}(x)_i) = x_i - \text{logSumExp}(x) - \left(\text{logSumExp}(x_{-i}) - \text{logSumExp}(x)\right) = x_i - \text{logSumExp}(x_{-i}) $$

[pawel-czyz]

Looks good to me!