From 2863ab8c237d5507e514a8c45a48b8a25fabd8f9 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sun, 16 Aug 2026 10:19:24 +0200 Subject: [PATCH 1/2] change rate parameter in Poisson to Float32 The VType of the rate parameter in Poisson is the same as the Type of the samples returned. For Poisson, that is Integer. But the rate should be a Float parameter. This commit decouples the two, --- .../dimwit/stats/IndependentDistributions.scala | 14 ++++++++++++-- .../scala/dimwit/stats/DistributionSuite.scala | 4 ++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala index 2132c49..970240b 100644 --- a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala +++ b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala @@ -216,12 +216,22 @@ class Exponential[T <: Tuple: Labels, V: IsFloating](val rate: Tensor[T, V]) ext ) stndExp / rate -class Poisson[T <: Tuple: Labels, V: IsInteger](val rate: Tensor[T, V]) extends IndependentDistribution[T, V]: +class Poisson[T <: Tuple: Labels, V: IsInteger](val rate: Tensor[T, Float32]) extends IndependentDistribution[T, V]: override def elementWiseLogProb(x: Tensor[T, V]): Tensor[T, LogProb] = liftPyTensor(jstats.poisson.logpmf(x.jaxValue, mu = rate.jaxValue)) override def sample(k: Random.Key): Tensor[T, V] = - liftPyTensor( + liftPyTensor(rate.shape, VType[V])( Jax.jrandom.poisson(k.jaxKey, lam = rate.jaxValue, shape = rate.shape.dimensions.toPythonProxy) ) + +object Poisson: + + /** Create a Poisson distribution from a rate tensor, sampling as Int32 counts */ + def apply[T <: Tuple: Labels](rate: Tensor[T, Float32]): Poisson[T, Int32] = + new Poisson[T, Int32](rate) + + /** Create a Poisson distribution from a rate tensor, sampling counts of the given integer type */ + def apply[T <: Tuple: Labels, V: IsInteger](rate: Tensor[T, Float32], vtype: VType[V]): Poisson[T, V] = + new Poisson[T, V](rate) diff --git a/core/src/test/scala/dimwit/stats/DistributionSuite.scala b/core/src/test/scala/dimwit/stats/DistributionSuite.scala index ab0de81..cb280bd 100644 --- a/core/src/test/scala/dimwit/stats/DistributionSuite.scala +++ b/core/src/test/scala/dimwit/stats/DistributionSuite.scala @@ -401,7 +401,7 @@ class DistributionSuite extends DimwitTest: describe("Poisson"): it("logProbs matches JAX"): - val rate = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1, 3, 10)) + val rate = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 3.0f, 10.0f)) val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1, 2, 8)) val dist = Poisson(rate) @@ -413,7 +413,7 @@ class DistributionSuite extends DimwitTest: it("sample means approximates rate"): val poisson = Poisson( - Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1, 5)) + Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 5.0f)) ) val key = Random.Key(42) val samples = key.splitvmap(Axis[Samples] -> 10000)(k => poisson.sample(k)) From 55273e439736b9f455b3815900156c9b32938c60 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sun, 16 Aug 2026 19:58:00 +0200 Subject: [PATCH 2/2] add separate vtype for rate --- .../stats/IndependentDistributions.scala | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala index 970240b..6929598 100644 --- a/core/src/main/scala/dimwit/stats/IndependentDistributions.scala +++ b/core/src/main/scala/dimwit/stats/IndependentDistributions.scala @@ -216,22 +216,22 @@ class Exponential[T <: Tuple: Labels, V: IsFloating](val rate: Tensor[T, V]) ext ) stndExp / rate -class Poisson[T <: Tuple: Labels, V: IsInteger](val rate: Tensor[T, Float32]) extends IndependentDistribution[T, V]: +class Poisson[T <: Tuple: Labels, VSample: IsInteger, VRate: IsFloating](val rate: Tensor[T, VRate]) extends IndependentDistribution[T, VSample]: - override def elementWiseLogProb(x: Tensor[T, V]): Tensor[T, LogProb] = + override def elementWiseLogProb(x: Tensor[T, VSample]): Tensor[T, LogProb] = liftPyTensor(jstats.poisson.logpmf(x.jaxValue, mu = rate.jaxValue)) - override def sample(k: Random.Key): Tensor[T, V] = - liftPyTensor(rate.shape, VType[V])( - Jax.jrandom.poisson(k.jaxKey, lam = rate.jaxValue, shape = rate.shape.dimensions.toPythonProxy) + override def sample(k: Random.Key): Tensor[T, VSample] = + liftPyTensor(rate.shape, VType[VSample])( + Jax.jrandom.poisson(k.jaxKey, lam = rate.jaxValue, shape = rate.shape.dimensions.toPythonProxy, dtype = VType[VSample].dtype.jaxType) ) object Poisson: /** Create a Poisson distribution from a rate tensor, sampling as Int32 counts */ - def apply[T <: Tuple: Labels](rate: Tensor[T, Float32]): Poisson[T, Int32] = - new Poisson[T, Int32](rate) + def apply[T <: Tuple: Labels, VRate: IsFloating](rate: Tensor[T, VRate]): Poisson[T, Int32, VRate] = + new Poisson[T, Int32, VRate](rate) /** Create a Poisson distribution from a rate tensor, sampling counts of the given integer type */ - def apply[T <: Tuple: Labels, V: IsInteger](rate: Tensor[T, Float32], vtype: VType[V]): Poisson[T, V] = - new Poisson[T, V](rate) + def apply[T <: Tuple: Labels, VSample: IsInteger, VRate: IsFloating](rate: Tensor[T, VRate], vtype: VType[VSample]): Poisson[T, VSample, VRate] = + new Poisson[T, VSample, VRate](rate)