Skip to content

feat(distributions): add LogNormal distribution#331

Open
Sumu004 wants to merge 1 commit into
google-deepmind:mainfrom
Sumu004:feat/add-lognormal-distribution
Open

feat(distributions): add LogNormal distribution#331
Sumu004 wants to merge 1 commit into
google-deepmind:mainfrom
Sumu004:feat/add-lognormal-distribution

Conversation

@Sumu004

@Sumu004 Sumu004 commented Jun 4, 2026

Copy link
Copy Markdown

What does this PR do?

Adds distrax.LogNormal(loc, scale) — the Log-Normal distribution on (0, ∞). If X ~ LogNormal(μ, σ), then log(X) ~ Normal(μ, σ).

The Log-Normal arises naturally in probabilistic ML wherever multiplicative noise is present: financial returns, rainfall, protein expression, and many Bayesian model priors. Unlike Beta/Gamma (bounded or semi-bounded), LogNormal covers the full positive real line with a simple reparameterizable sampler.

Implemented

Method Formula
log_prob(x) Normal(μ,σ).log_prob(log x) − log x
mean() exp(μ + σ²/2)
variance() (exp(σ²)−1) · exp(2μ+σ²)
mode() exp(μ − σ²)
median() exp(μ)
entropy() μ + ½ + log(σ) + ½ log(2π) (exact)
kl_divergence Analytic when other is LogNormal (= Normal KL in log-space)
sample() exp(Normal(μ,σ).sample()) — reparameterizable

Tests

17 tests in log_normal_test.py: log_prob vs TFP and scipy, mean / variance / mode / median / entropy vs TFP, KL divergence (self=0, analytic vs TFP), sample shape, positivity, moments, JIT compatibility.

Adds distrax.LogNormal(loc, scale) — the Log-Normal distribution on
(0, ∞), where log(X) ~ Normal(loc, scale).

The Log-Normal arises naturally wherever multiplicative noise is
present: financial returns, rainfall intensity, protein expression
levels, and many Bayesian model priors.

Implements:
  - log_prob(x) = Normal(mu, sigma).log_prob(log x) − log x
  - mean()     = exp(μ + σ²/2)
  - variance() = (exp(σ²)−1) * exp(2μ+σ²)
  - mode()     = exp(μ − σ²)
  - median()   = exp(μ)
  - entropy()  = μ + ½ + log(σ) + ½ log(2π)  (exact closed form)
  - kl_divergence(other): delegates to the underlying Normal KL when
    other is also LogNormal (analytic, equals Normal KL in log-space)

Adds distrax.LogNormal to the public __init__.py and 17 tests covering
log_prob against both TFP and scipy, mean/variance/mode/median/entropy
against TFP, KL divergence, sample shape/positivity/moments, and JIT.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant