A Jax implementation of word2vec's skip-gram model with negative sampling as described in Mikolov et al., 2013