distrax_multivariate_normal_log_prob_and_sampling_with_jax.py

1import distrax
2import jax
3import jax.numpy as jnp
4
5# Define the parameters of a multivariate Normal distribution.
6mu = jnp.zeros(2)
7sigma = jnp.ones(2)
8
9# Initialize the distribution.
10dist = distrax.MultivariateNormalDiag(loc=mu, scale_diag=sigma)
11
12# Calculate the log-density of a point.
13lp = dist.log_prob(jnp.zeros(2))
14print(f"Log probability: {lp}")
15
16# Sample from the distribution.
17seed = jax.random.PRNGKey(42)
18samples = dist.sample(seed=seed, sample_shape=(5,))
19print(f"Samples:\n{samples}")