Blog

🧠 Introduction: What Is torch.distributions?

Probabilistic modeling is at the core of many machine learning and deep learning algorithms—from variational autoencoders (VAEs) to Bayesian inference. PyTorch offers a powerful, flexible module to handle these needs: torch.distributions.

torch.distributions is a PyTorch subpackage that provides a rich set of probability distributions and probabilistic tools. It allows you to define, sample from, compute probabilities, and manipulate distributions—essential for tasks in uncertainty modeling, generative models, and reinforcement learning.

---

🛠️ Code Examples: Creating and Sampling Distributions

Here are some practical examples of how to use torch.distributions.

1. Create a Normal (Gaussian) Distribution

pythonCopyEditimport torch
from torch.distributions import Normal

# Define a normal distribution with mean=0 and std=1
normal_dist = Normal(loc=0.0, scale=1.0)

# Sample from the distribution
sample = normal_dist.sample((5,))
print("Samples:", sample)

---

2. Log Probability (Likelihood)

pythonCopyEditlog_prob = normal_dist.log_prob(torch.tensor(0.0))
print("Log Probability of 0.0:", log_prob.item())

---

3. Other Distributions

pythonCopyEditfrom torch.distributions import Bernoulli, Categorical

# Bernoulli (binary outcomes)
bernoulli_dist = Bernoulli(probs=0.7)
print("Bernoulli sample:", bernoulli_dist.sample())

# Categorical (multi-class)
categorical = Categorical(probs=torch.tensor([0.1, 0.3, 0.6]))
print("Categorical sample:", categorical.sample())

---

4. Reshape (Batch) Distributions

pythonCopyEdit# Create a batched Normal distribution
batched_dist = Normal(torch.zeros(3), torch.ones(3))
samples = batched_dist.sample()
print("Batched Normal sample:", samples)

---

🔄 Commonly Used Classes & Methods in torch.distributions

Class	Description
Normal	Gaussian distribution with mean (loc) and standard deviation (scale)
Bernoulli	Binary discrete distribution with probability p
Categorical	Multi-class distribution defined by class probabilities
MultivariateNormal	Multidimensional Gaussian
Poisson	Distribution of rare events
Beta, Gamma, Dirichlet	Common for Bayesian models

Method	Purpose
.sample()	Generates a random sample from the distribution
.log_prob(x)	Computes log probability of input x
.rsample()	Reparameterized sample for gradient flow (used in VAEs)
.entropy()	Measures uncertainty in the distribution
.mean, .variance, .stddev	Statistical properties

---

❗ Errors and Debugging Tips

❌ Error 1: ValueError: Expected parameters to satisfy constraints

Cause: Your distribution’s parameters violate mathematical rules (e.g., negative standard deviation).

Fix:

pythonCopyEditfrom torch.distributions import constraints

# Use positive scale
Normal(loc=0.0, scale=torch.tensor(1.0))  # Valid

---

❌ Error 2: RuntimeError: The size of tensor a (x) must match tensor b (y)

Cause: Mismatched tensor shapes during operations like .log_prob().

Fix: Ensure inputs and distribution parameters are broadcastable:

pythonCopyEditNormal(torch.zeros(3), torch.ones(3)).log_prob(torch.tensor([0.1, 0.2, 0.3]))

---

❌ Error 3: Gradients not flowing in VAEs

Fix: Use .rsample() instead of .sample() for reparameterization:

pythonCopyEditsample = normal_dist.rsample()

This ensures gradients can propagate during training.

---

❌ Error 4: Unexpected behavior in batches

Tip: Always check the shape of your distribution and samples:

pythonCopyEditdist = Normal(torch.zeros(2), torch.ones(2))
print(dist.batch_shape)  # Output: torch.Size([2])

---

✅ People Also Ask (FAQ)

🔹 What Is torch.distributions Used For?

torch.distributions is used for defining and sampling from probability distributions in PyTorch. It supports statistical computations like log probabilities, entropy, and reparameterized sampling for use in generative models like VAEs, reinforcement learning, and uncertainty modeling.

---

🔹 What Is the Difference Between .sample() and .rsample()?

• .sample(): Generates samples without tracking gradients.
• .rsample(): Enables reparameterized sampling, allowing gradients to flow through stochastic nodes. Useful in training VAEs or probabilistic models with backpropagation.

---

🔹 How Do You Use Categorical in PyTorch?

pythonCopyEditfrom torch.distributions import Categorical
probs = torch.tensor([0.2, 0.3, 0.5])
cat = Categorical(probs=probs)
sample = cat.sample()

This is often used to sample discrete actions in reinforcement learning.

---

🔹 How Do You Compute the Log-Likelihood of a Sample?

Use .log_prob(x):

pythonCopyEditdist = Normal(0.0, 1.0)
x = torch.tensor(0.5)
log_likelihood = dist.log_prob(x)

This is essential in Bayesian inference and ELBO computation.

---

🔹 Can I Define Custom Distributions?

Yes! You can create your own distribution by subclassing torch.distributions.Distribution and implementing required methods like .sample() and .log_prob().

---

📌 Final Thoughts

The torch.distributions module is a cornerstone for anyone building probabilistic models in PyTorch. Whether you’re working on generative models, reinforcement learning, or uncertainty quantification, this library equips you with the statistical tools you need.

Here’s a recap of what makes torch.distributions powerful:

• 🚀 Wide range of built-in distributions
• ✅ Reparameterized sampling with .rsample()
• 🔍 Easy computation of log-probabilities, entropy, and other metrics
• 🔄 Supports batched operations for efficient parallelism
• 🔧 Can be extended to build custom distributions

Understanding and mastering torch.distributions opens the door to advanced techniques like variational inference, probabilistic programming, and more.