Replaced seed with key in add_noise and noisy_sgd

[Tomas542]

Replaced seed in add_noise with key in favor of jax.random-like style.
Added (duplicated) example with add_noise from noisy_sgd.
Changed seed to key in noisy_sgd. Imported chex for annotation purpose in _alias.py

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[rdyro]

These changes look good! Can you sign the google CLA to be able to run unit tests?

I like the change of placing the key first in the argument list. Can you take a look at these files as well:

• optax/optax/_src/utils.py

  Line 112 in 63cdeb4

  def sample(self, shape: Sequence[int], seed: chex.PRNGKey) -> chex.Array:

• optax/optax/perturbations/_make_pert.py

  Line 38 in 63cdeb4

  seed: chex.PRNGKey,

• optax/optax/contrib/_privacy.py

  Line 36 in 63cdeb4

  l2_norm_clip: float, noise_multiplier: float, seed: int

It'd be good to rename the argument to key from seed in all these places and move the key argument to the beginning of the argument list - this might require touching up a couple of tests.

[Tomas542]

I signed the CLA yesterday and went to bed.
It's a great idea, but I think I'll do it this weekend.
I also have a question (sorry if it's stupid, I've never done Pull Requests before): if I make a new commit with these changes to my branch - do they automatically go into this PR? Or is there some button I have to press?

[rdyro]

Awesome

if I make a new commit with these changes to my branch - do they automatically go into this PR? Or is there some button I have to press?

Not at all, thanks for the contribution! Exactly, any additional commits in your branch will be added to the PR here, in the end, you can also squish your commits into one before merging with something like:

$ git rebase -i main
$ # change all but one commit to "squish" leaving one of them (probably the first one) as "pick"
$ # if everything went well
$ git push -f (to overwrite history, which is necessary when squishing commits)

[Tomas542]

Ok, and one more question - there are a lot of things that will change. Maybe we'll add key as an optional argument with an DeprecationWarning if seed is set, like Warning: argument seed will be removed (or replaced in cases where seed already PRNGKey) in the next release. Use key instead. But in this case key will move from the first argument, and we will have to make it of type Optional for now. And in the next-next major release I will remove/replace seed with key.
Or is it better to change everything now? What do you think?

[vroulet]

Just a small comment about the signature of the optimizer. Otherwise I fully agree with Robert that a full pass would be great. Could be in a separate PR if needed.

[vroulet]

I would still keep learning_rate first as in other optimizers for this function, and I would give a default value for the key for convenience, like key=jax.random.PNRGKey(0).
Sometimes one wants to benchmark quickly optimizers and it seems easier if they all have a similar signature, at least for the first argument.
On the other hand, for the transform, (add_noise) and other functions Robert pointed out, I agree that it would make sense to have the key first.