Re:Show HN: Word2vec Algorithm in ~100sloc with NumPy
by sdenton4 on Thursday June 01, 2023 @11:16AM
(#36152321)
Neat! It /should/ be pretty easy to convert this to Jax, as a big party of the Jax API is just reimplementing the Numpy API.
The two changes I can imagine are: a) handling randomness requires more explicit rng estate handling, and b) wrapping the inner loop in a function so that it can be JIT-compiled.
The advantage to doing this is that you can then run your code on a GPU, getting you most of the way to 'blazingly fast.'
The two changes I can imagine are: a) handling randomness requires more explicit rng estate handling, and b) wrapping the inner loop in a function so that it can be JIT-compiled.
The advantage to doing this is that you can then run your code on a GPU, getting you most of the way to 'blazingly fast.'