From 39a4c0366188c4ef9fb13a05cb51a038ce4cdc85 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Mon, 2 Nov 2026 17:20:15 +0100 Subject: [PATCH] Add Gamma distribution implementation and corresponding tests --- .../stats/IndependentDistributions.scala | 31 ++++++++++ .../dimwit/stats/DistributionSuite.scala | 58 +++++++++++++++++++ 2 files changed, 89 insertions(+) diff --git a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala index 4f37980..84c19f5 100644 --- a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala +++ b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala @@ -243,6 +243,37 @@ class Exponential[T <: Tuple: Labels, V: IsFloating](val rate: Tensor[T, V]) ext ) stndExp / rate +/** Gamma distribution with `concentration` (the shape parameter, often written alpha or k) and `rate` (often written beta). + * + * The parameter is called concentration rather than shape, as in Dirichlet, to avoid a clash with a tensor's `shape`. + * Mean is concentration / rate, variance concentration / rate^2. Gamma(1, rate) is Exponential(rate). + */ +class Gamma[T <: Tuple: Labels, V: IsFloating](val concentration: Tensor[T, V], val rate: Tensor[T, V]) extends IndependentDistribution[T, V]: + + private lazy val scale = Tensor.like(rate).fill(1f) / rate + + override def elementWiseLogProb(x: Tensor[T, V]): Tensor[T, LogProb] = + LogProb.fromPy(jstats.gamma.logpdf(x.jaxValue, a = concentration.jaxValue, scale = scale.jaxValue)) + + override def sample(k: Random.Key): Tensor[T, V] = + // jax.random.gamma samples with rate 1, so divide by the rate + val standardGamma = liftPyTensor(concentration.shape, VType[V])( + Jax.jrandom.gamma(k.jaxKey, a = concentration.jaxValue, shape = concentration.shape.dimensions.toPythonProxy, dtype = VType[V].dtype.jaxType) + ) + standardGamma / rate + +object Gamma: + + /** Create a Gamma distribution from concentration and rate tensors */ + def apply[T <: Tuple: Labels, V: IsFloating](concentration: Tensor[T, V], rate: Tensor[T, V]): Gamma[T, V] = + require(concentration.shape.dimensions == rate.shape.dimensions, "concentration and rate must have the same dimensions") + new Gamma(concentration, rate) + + /** Create a Gamma distribution from concentration and scale tensors, where scale = 1 / rate */ + def fromScale[T <: Tuple: Labels, V: IsFloating](concentration: Tensor[T, V], scale: Tensor[T, V]): Gamma[T, V] = + require(concentration.shape.dimensions == scale.shape.dimensions, "concentration and scale must have the same dimensions") + new Gamma(concentration, Tensor.like(scale).fill(1f) / scale) + class Poisson[T <: Tuple: Labels, VSample: IsInteger, VRate: IsFloating](val rate: Tensor[T, VRate]) extends IndependentDistribution[T, VSample]: override def elementWiseLogProb(x: Tensor[T, VSample]): Tensor[T, LogProb] = diff --git a/core/src/test/scala/dimwit/stats/DistributionSuite.scala b/core/src/test/scala/dimwit/stats/DistributionSuite.scala index 09dbaeb..420e5fa 100644 --- a/core/src/test/scala/dimwit/stats/DistributionSuite.scala +++ b/core/src/test/scala/dimwit/stats/DistributionSuite.scala @@ -567,6 +567,63 @@ class DistributionSuite extends DimwitTest: val expectedMeans = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(2.0f, 0.5f)) sampleMeans should approxEqual(expectedMeans, 0.1f) + describe("Gamma"): + it("logProbs matches JAX"): + val concentration = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 2.0f, 5.0f)) + val rate = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 0.5f, 4.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.3f, 2.5f, 1.2f)) + + val dist = Gamma(concentration, rate) + val scalaLogProbs = dist.elementWiseLogProb(x) + val scale = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 2.0f, 0.25f)) + val jaxLogProbs = liftPyTensor1(Axis[A], VType[Float32])( + jstats.gamma.logpdf(x.jaxValue, a = concentration.jaxValue, scale = scale.jaxValue) + ) + scalaLogProbs.asFloat should approxEqual(jaxLogProbs) + + it("fromScale is the same as the rate parameterisation"): + val concentration = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(2.0f, 3.0f)) + val scale = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(2.0f, 0.5f)) + val rate = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.5f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 1.5f)) + + Gamma.fromScale(concentration, scale).elementWiseLogProb(x).asFloat should approxEqual(Gamma(concentration, rate).elementWiseLogProb(x).asFloat) + + it("Gamma(1, rate) is Exponential(rate)"): + val rate = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.5f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.7f, 1.3f)) + val gamma = Gamma(Tensor.like(rate).fill(1.0f), rate) + + gamma.elementWiseLogProb(x).asFloat should approxEqual(Exponential(rate).elementWiseLogProb(x).asFloat) + + it("logProb is -inf for x < 0"): + val concentration = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(2.0f, 3.0f)) + val rate = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(-1.0f, -0.1f)) + + val dist = Gamma(concentration, rate) + val logProbs = dist.elementWiseLogProb(x) + logProbs.asFloat.toArray.foreach(v => v should be(Float.NegativeInfinity)) + + it("sample means approximates concentration / rate"): + val gamma = Gamma( + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(2.0f, 9.0f)), + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(4.0f, 3.0f)) + ) + val key = Random.Key(42) + val samples = key.splitvmap(Axis[Samples] -> 10000)(k => gamma.sample(k)) + val sampleMeans = samples.mean(Axis[Samples]) + val expectedMeans = gamma.concentration / gamma.rate + sampleMeans should approxEqual(expectedMeans, 0.1f) + + it("works with scalar parameters"): + val gamma = Gamma(Tensor0(2.0f), Tensor0(4.0f)) + val key = Random.Key(7) + val samples = key.splitvmap(Axis[Samples] -> 10000)(k => gamma.sample(k)) + samples.mean.item shouldBe 0.5f +- 0.05f + val jaxLogProb = liftPyTensor0(VType[Float32])(jstats.gamma.logpdf(1.0f, a = 2.0f, scale = 0.25f)) + gamma.logProb(Tensor0(1.0f)).asFloat should approxEqual(jaxLogProb) + describe("Poisson"): it("logProbs matches JAX"): val rate = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 3.0f, 10.0f)) @@ -618,6 +675,7 @@ class DistributionSuite extends DimwitTest: checkTypes("StudentT", StudentT(Tensor0(3.0f).asType(VType[Float16]), loc, scale), x) checkTypes("Beta", Beta(f16(2.0f, 3.0f), f16(2.0f, 1.0f)), f16(0.3f, 0.6f)) checkTypes("Exponential", Exponential(f16(1.0f, 2.0f)), x) + checkTypes("Gamma", Gamma(f16(2.0f, 3.0f), f16(1.0f, 2.0f)), x) checkTypes("Dirichlet", Dirichlet(f16(1.0f, 2.0f)), f16(0.3f, 0.7f)) it("TruncatedNormal samples Float16"):