elixir-nx / nx

Multi-dimensional arrays (tensors) and numerical definitions for Elixir
2.66k stars 194 forks source link

More docs and seed usage guidance needed on `Nx.Random` #1515

Closed joshprice closed 1 month ago

joshprice commented 4 months ago

The docs for Nx.Random https://hexdocs.pm/nx/Nx.Random.html describe the usage of Nx.Random.key/1 (https://hexdocs.pm/nx/Nx.Random.html#key/1) but only have examples of using literal integers as the seed.

Setting the key to a literal integer seed ensures that the results of any random function call will be the same. This probably isn't typically what a user would want (perhaps you do in a test case or if you're trying to reproduce an error or something) - as per the discussion for Nx.Random.

I didn't see any other docs or guides on this. Would be good to have a more realistic example and some guidance around how to manage keys so that you get properly random numbers.

Perhaps this just pushes the seed problem another just one more function away, but I'm imagining that you might want Nx.Random.key/0 which returns something like :rand.uniform(1_000_000_000) or :rand.uniform(1_000_000) * System.system_time(:nanosecond).

josevalim commented 4 months ago

System.os_time is a good value but you can also use :rand.uniform(...).

polvalente commented 3 months ago

@joshprice a PR improving the moduledoc would probably be nice. Nx.Random.key/0 doesn't seem to be the right route to me. Take for example this code:

# Let's take key/0 as:
def key do
  Nx.Random.key(:rand.uniform(1_000_000))
end

:rand.seed(:default, 1)

# Process 1:
{uniform1, key1} = Nx.Random.uniform(key())

# Process 2:
{uniform2, key2} = Nx.Random.uniform(key())

This code returns different results for uniform1 and uniform2 solely depending on which process calls key() first. This would seem desirable for someone that wants random results, but also takes away from the main purpose of Nx.Random which is to provide a stateless PRNG interface.

That being said, if we're having this discussion here, it means that we can improve documentation. The moduledoc does mention the correct workflow: pass the newly returned key as the key to the next invokation. Perhaps you could take a stab at wording it in a better way or with a more extensive example? Maybe add part of this discussion there?