Skip to content
← Learn JAX

JAX Stochasticity

JAX

Distributions, reparameterization, sampling techniques, MCMC, gradient estimators. Randomness as a controllable resource.

This track uses JAX. You can make it your practice mode to browse related content across the site.

25 exercises

Explore freely. Sign in to keep your progress.

Start track

Your learning sequence

  1. 1 Not solved yet. Uniform Sampling Practice · Easy
  2. 2 Not solved yet. Normal Sampling with mean and std Practice · Easy
  3. 3 Not solved yet. Bernoulli Mask Sampling Practice · Easy
  4. 4 Not solved yet. Categorical Sampling Practice · Medium
  5. 5 Not solved yet. Reparameterization Trick: Gaussian Practice · Medium
  6. 6 Not solved yet. Gumbel-Softmax Practice · Hard
  7. 7 Not solved yet. Dirichlet Sampling Practice · Medium
  8. 8 Not solved yet. Decoding Temperature Practice · Easy
  9. 9 Not solved yet. Top-k Logit Masking Practice · Medium
  10. 10 Not solved yet. Nucleus (Top-p) Masking Practice · Hard
  11. 11 Not solved yet. Gumbel-Max Categorical Sampling Practice · Medium
  12. 12 Not solved yet. Metropolis-Hastings Step Practice · Medium
  13. 13 Not solved yet. Log Acceptance Ratio Practice · Easy
  14. 14 Not solved yet. Leapfrog with a Constant Force Practice · Hard
  15. 15 Not solved yet. REINFORCE Gradient Estimator Practice · Medium
  16. 16 Not solved yet. Reparameterization Gradient Practice · Medium
  17. 17 Not solved yet. REINFORCE with Baseline Practice · Hard
  18. 18 Not solved yet. Batched Sampling via vmap Practice · Medium
  19. 19 Not solved yet. Random Walk via lax.scan Practice · Hard
  20. 20 Not solved yet. Importance Sampling Practice · Hard
  21. 21 Not solved yet. Beta Distribution Sampling Practice · Easy
  22. 22 Not solved yet. Multivariate Normal Sampling Practice · Medium
  23. 23 Not solved yet. Random Permutation Practice · Easy
  24. 24 Not solved yet. Poisson Sampling Practice · Medium
  25. 25 Not solved yet. Gaussian Reconstruction and KL Practice · Hard

Check yourself

4 questions · not saved while you're signed out

Use these questions to check your understanding and compare explanations. Your score does not affect track completion.

0 / 4 correct

With key-reuse checking disabled, dropout calls in multiple layers and steps reuse the same key, shape, and probability. What masks do these calls produce?

Choose one answer to reveal the explanations. Answers are kept only while this page is open. 

key = jax.random.key(0)
shape = (8,)
# Each layer and step repeats this same call:
jax.random.bernoulli(key, 0.5, shape)
Question 1 of 4