Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Register a random module in jax when registering jax modules #16

Open
to24toro opened this issue May 24, 2023 · 0 comments
Open

Register a random module in jax when registering jax modules #16

to24toro opened this issue May 24, 2023 · 0 comments
Labels
enhancement New feature or request

Comments

@to24toro
Copy link
Collaborator

to24toro commented May 24, 2023

What is the expected behavior?

If you want to use random.seed, random.unform, and so on in jax, you have to use from jax.random import PRNGKey and how to use the random module is different from the numpy one.
So I would like to implement functions for random module in jax, make it so simpler and enable us to use like numpy random when registering jax modules.

I would like to use like

jax.random.seed(42)
jax.random.normal() # this produce the same results as long as a certain random seed is specified.

# As a default, jax.random.normal is provided but this requires PRNG key argument.
@to24toro to24toro added the enhancement New feature or request label May 24, 2023
@to24toro to24toro changed the title Registering a module like random module in jax as default Registering a random module in jax when registering jax modules May 24, 2023
@to24toro to24toro changed the title Registering a random module in jax when registering jax modules Register a random module in jax when registering jax modules May 24, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

1 participant