A bit petty, but the first example is an unstable softmax implementation:
defn softmax(t) do
Nx.exp(t) / Nx.sum(Nx.exp(t))
end
See https://ogunlao.github.io/2020/04/26/you_dont_really_know_so... etc.