nx

nx copied to clipboard

• Implementation of stateless PRNG, based on splittable keys and the Threefry algorithm
• Based on the JAX PRNG implementation and yield the same results
• At the moment includes only randint and uniform generators

#333

I've implemented most of the suggestions, some of them are still a little clunky since we don't really have support for 64bits. Probably missed some things, will look for them later.

Hi @john-sonz! I am checking on the status of this PR. Anything pending? :) Can we help? TY! :heart:

Hi @john-sonz! I am checking on the status of this PR. Anything pending? :) Can we help? TY! ❤️

This PR should be good to go after the code formatting ( not counting test errors from linalg :') ) and we'll hopefully start working on vmap. @john-sonz mostly covered batching in linalg functions.

Thanks for this work! I have opened a corollary to this PR in Axon here: https://github.com/elixir-nx/axon/pull/350

I noticed that right now Nx.Random.uniform does not correctly preserve the requested type, you may have to cast at the end or cast some of the intermediate values

Excellent @john-sonz and @VenomSolo! I have done my review and it is mostly minor details to clean up the code and docs. I have also started working on the basic foundation for vmap and I hope to have something to share soon. :)

Made an adjustment to reduce slicing suggested earlier by @polvalente but somehow glossed over.

Amazing work! :green_heart: :blue_heart: :purple_heart: :yellow_heart: :heart: