add_noise should take key as KeyArrayLike, not seed as int

[Tomas542]

I think, that seed in add_noise should be replaced with key to follow jax.random and flax.nnx.Modules, where keys are used instead of seeds. It doesn't seem hard to do, like replace seed in args, in docstring, and then in AddNoiseState init (AddNoiseState by itself doesn't need to change as far as I can see).

[Tomas542]

also key init in AddNoiseState should be changed with new jax.random.key in favor of JEP

[rdyro]

You're right! Do you want to contribute a PR fixing this?

[Tomas542]

I think I can, but I've never done this before. Should I check contributor's guide on official cite?

UPD.
Ok, I've read official development guidelines, is this the right way to do things:

• For Python code I need just to git clone optax repository and change code I've mentioned earlier (also I think it is great idea to add example).
• run bash test.sh to check whether it is ok.
  And that's it? Documentation will update automatically?

[rdyro]

That's right. If you want to pursue it, it's a valuable contribution.

You can fork the repository, make a new branch in your fork, add the changes (to both code and comments) and update tests (probably no need to write new tests). If you point your browser to this repo, github should offer to make a pull request from your fork.

Feel free to do it incrementally, you can make the PR almost immediately and any additional commits in the branch in your fork will automatically show up in the PR.

I can review your code once you make the PR

[Tomas542]

Done! But currently I have no option to run test.sh on my local machine. Can I create PR? Or better to set virtual machine with Linux for test?

[Tomas542]

Also I additionally changed noisy_sgd, cause it's uses add_noise and added (duplicated) example of use from noisy_sgd to add_noise

[rdyro]

Cool! Let's continue this discussion on the PR, can you make sure you create it for deepmind/optax? I can't see it at the moment

[Tomas542]

here #1138