How to set the mean and variance of normal distribution?

[Dsantra92]

How to generate normal distribution of given mean and std. deviation in Jax? To be more explicit, I am looking for a function corresponding to the torch function torch.normal(mu, sigma, shape).

[jakevdp]

You can do this by generating standard normal numbers with jax.random.normal, then multiply by sigma, and add mu:

from jax import random

mu = 1.0
sigma = 3.0
key = random.PRNGKey(1701)
x = mu + sigma * random.normal(key, shape=(10000,))

print(x.mean(), x.std())
# 0.9974507 2.9847035

(note that sigma here represents the standard deviation rather than the variance; I believe this is the same convention used in torch.normal).

[Dsantra92]

Thank you, your solution works. I was curious as to why the function does not have an option for manual values, it would be really nice if the normal() function has an inbuilt support for user-defined parameters.

ps- Yes it is indeed std. deviation. I have updated the discussion thread.

[cisprague]

Would also be nice if jax.random.multivariate_normal accepted cov of different shapes based on the shape of mean.shape == (..., n). If cov.shape == (...) use scalar standard deviation method. If (..., n) use diagonal covariance method. If (..., n, n) use one of ('eigh', 'svd', 'cholesky').

[jakevdp]

Generalizing cov that way leads to ambiguities. Say for example mean has shape (3, 3, 3), and you pass cov of shape (3, 3). Is that an unbatched 2D covariance matrix, a singly-batched 1D diagonal, or a doubly-batched scalar? Any of those options would be valid in terms of your proposed specification. The benefit of the current approach, requiring explicit 2D covariance, is that it leaves no room for such ambiguities.

[cisprague]

Thanks for the answer, that makes sense. Another question: is there a way to let jax.random.multivariate_normal know we are passing a diagonal cov matrix, to avoid using eigh, svd, or cholesky? If not, then adding a method called diagonal could be nice.

[jakevdp]

That's a good idea – could you open a new issue with a feature request? I'm worried this idea will get lost in the shuffle if we leave it in a deep comment on a 3-year-old issue 😁