Pass key instead of seed to algorithms with randomization

[msluszniak]

If we want this functionality I can implement this.

[josevalim]

Let's try to apply this to one model and see how it looks like? Perhaps we can allow one of key or seed to be given?

[polvalente]

I'd rather just allow key for uniformity, because that's our equivalent of seed in Nx land

[josevalim]

@polvalente but it makes the API harder to use for everyone, doesn't it?

[polvalente]

I don't think it does. If the key is optional, it's just what we have now.

If we accept a seed, it's much harder to use than a key in terms of defn usage because the natural System.system_time call isn't available as well.

I don't see any normal use case where a seed is really needed instead of a key.

Also, we need to either return the key or teach people to split the key beforehand.

[josevalim]

Also, we need to either return the key or teach people to split the key beforehand.

That's exactly the point of how it makes it harder. For example, if you write a notebook, you may want to hardcode the seed in the notebook. But now we have to tell them to create a key before hand and split it?

[polvalente]

Split is only needed if you don't return the key, which I think is the preferred approach.

[josevalim]

Right, so we either have to return the key in all APIs or the user has to split before hand instead of passing a seed. All of those look like more work to me than just passing an integer. :) Plus, we may end-up generating the key outside of a defn, which will affect performance. So I think supporting both is the best of both worlds?

[polvalente]

How does generating the key outside of defn affect performance?

[josevalim]

@polvalente all of the Nx code to generate the key will be executed outside of defn and therefore it won't be compiled?

[josevalim]

Nah, ignore me. key is very cheap to compute. :)

[josevalim]

So I guess for the cases we don't need to split, we just need to replace seed: 42 by key: Nx.Random.key(42)? If that's the case, rolling with :key is 100% fine by me.

[polvalente]

key must be a tensor input, though, because it is split and carried over multiple computations

[josevalim]

Can you please expand, so we are in the same page?

[polvalente]

Imagine we're using one of Scholar's algorithms that needs to receive a random key in the middle of some defn code.

If the previous code already used the key, it'll have a computation graph associated with it, even if it's just a sequence of splits.

And from my understanding, passing that as an option would be bad for the Nx compiler.

[josevalim]

And from my understanding, passing that as an option would be bad for the Nx compiler.

We handle this today in Scholar by defining all fit functions as transform and we convert options into inputs.

[polvalente]

Nice, so I'm ok with having key as option!

[josevalim]

Alright, so :key as an option it is!