Add Gamma distribution implementation and corresponding tests

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] =

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"):