distrax_jax_normal_distribution_sampling_and_log_prob.py

1import distrax
2import jax
3import jax.numpy as jnp
4
5# Create a normal distribution with mean 0 and standard deviation 1.
6dist = distrax.Normal(loc=0.0, scale=1.0)
7
8# Sample from the distribution.
9key = jax.random.PRNGKey(42)
10samples = dist.sample(seed=key, sample_shape=(5,))
11
12# Compute the log-probability of the samples.
13log_probs = dist.log_prob(samples)
14
15print(f"Samples: {samples}")
16print(f"Log-probabilities: {log_probs}")