diff --git a/core/src/main/scala/dimwit/stats/UnivariateDistributions.scala b/core/src/main/scala/dimwit/stats/UnivariateDistributions.scala index 06c5b25..6185086 100644 --- a/core/src/main/scala/dimwit/stats/UnivariateDistributions.scala +++ b/core/src/main/scala/dimwit/stats/UnivariateDistributions.scala @@ -19,7 +19,10 @@ class Categorical[L: Label](val probs: Tensor1[L, Prob]) extends UnivariateDistr private val logProbs: Tensor1[L, LogProb] = probs.log override def logProb(x: Tensor0[Int32]): Tensor0[LogProb] = - liftPyTensor(logProbs.jaxValue.__getitem__(x.jaxValue)) + // JAX clamps out-of-range indices (and wraps negative ones), so values outside 0..K-1 are masked explicitly + val indexed: Tensor0[LogProb] = liftPyTensor(logProbs.jaxValue.__getitem__(x.jaxValue)) + val inRange = (x >= Tensor0(0)) and (x < Tensor0(probs.shape.dimensions(0))) + where(inRange, indexed, LogProb(Tensor0(Float.NegativeInfinity))) override def sample(key: Key): Tensor0[Int32] = liftPyTensor(Jax.jrandom.categorical(key.jaxKey, logProbs.jaxValue)) diff --git a/core/src/test/scala/dimwit/stats/DistributionSuite.scala b/core/src/test/scala/dimwit/stats/DistributionSuite.scala index ae30f23..09dbaeb 100644 --- a/core/src/test/scala/dimwit/stats/DistributionSuite.scala +++ b/core/src/test/scala/dimwit/stats/DistributionSuite.scala @@ -63,6 +63,17 @@ class DistributionSuite extends DimwitTest: ) scalaLogProbs.asFloat should approxEqual(jaxLogProbs) + it("logProb is -inf outside [low, high]"): + val loc = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 0.0f)) + val scale = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 1.0f)) + val low = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(-1.0f, -1.0f)) + val high = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 1.0f)) + val x = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(-2.0f, 2.0f)) + + val dist = TruncatedNormal(loc, scale, low, high) + val logProbs = dist.elementWiseLogProb(x) + logProbs.asFloat.toArray.foreach(v => v should be(Float.NegativeInfinity)) + it("sample means approximates means"): val loc = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 1.0f)) val scale = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 0.5f)) @@ -99,6 +110,15 @@ class DistributionSuite extends DimwitTest: ) scalaLogProbs.asFloat should approxEqual(jaxLogProbs) + it("logProb is -inf outside [low, high]"): + val low = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 0.0f)) + val high = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 1.0f)) + val x = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(-0.5f, 1.5f)) + + val dist = Uniform(low, high) + val logProbs = dist.elementWiseLogProb(x) + logProbs.asFloat.toArray.foreach(v => v should be(Float.NegativeInfinity)) + it("sample means approximates means"): val uniform = Uniform( Tensor(Shape(Axis[A] -> 2)).fromArray(Array(-1.0f, 0.0f)), @@ -190,6 +210,15 @@ class DistributionSuite extends DimwitTest: ) scalaLogProbs.asFloat should approxEqual(jaxLogProbs) + it("logProb is -inf outside [0, n]"): + val n = Tensor0(10) + val probs = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.3f, 0.5f)) + val x = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(-1, 11)) + + val dist = Binomial(n, Prob(probs)) + val logProbs = dist.elementWiseLogProb(x) + logProbs.asFloat.toArray.foreach(v => v should be(Float.NegativeInfinity)) + it("sample means approximates n*p"): val n = Tensor0(20) val binomial = Binomial( @@ -300,6 +329,15 @@ class DistributionSuite extends DimwitTest: ) +! math.log(2.0).toFloat scalaLogProbs.asFloat should approxEqual(expectedLogProbs) + it("logProb is -inf for x < loc"): + val loc = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 1.0f, 1.0f)) + val scale = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1.0f, 1.0f, 1.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.0f, 0.5f, 0.9f)) + + val dist = HalfNormal(loc, scale) + val logProbs = dist.elementWiseLogProb(x) + logProbs.asFloat.toArray.foreach(v => v should be(Float.NegativeInfinity)) + it("sample means approximates expected means"): val halfNormal = HalfNormal( Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.0f, 0.0f)), @@ -382,6 +420,13 @@ class DistributionSuite extends DimwitTest: ) scalaLogProb.asFloat should approxEqual(jaxLogProb) + it("logProb is -inf off the simplex"): + val concentration = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(2.0f, 3.0f, 5.0f)) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.5f, 0.6f, 0.2f)) // sums to 1.3 + + val dist = Dirichlet(concentration) + dist.logProb(x).asFloat.item should be(Float.NegativeInfinity) + it("sample mean approximates expected mean"): val concentration = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(2.0f, 5.0f, 3.0f)) val dirichlet = Dirichlet(concentration) @@ -407,6 +452,13 @@ class DistributionSuite extends DimwitTest: ) scalaLogProb.asFloat should approxEqual(jaxLogProb) + it("logProb is -inf for counts that do not sum to n"): + val probs = Prob(Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.2f, 0.5f, 0.3f))) + val x = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(1, 1, 1)) + + val dist = Multinomial[A](Tensor0(5), probs) + dist.logProb(x).asFloat.item should be(Float.NegativeInfinity) + it("sample mean approximates expected counts"): val probsFloat = Tensor(Shape(Axis[A] -> 3)).fromArray(Array(0.2f, 0.5f, 0.3f)) val probs = Prob(probsFloat) @@ -429,6 +481,14 @@ class DistributionSuite extends DimwitTest: val expectedLogProb = Tensor0(math.log(0.3f).toFloat) scalaLogProb.asFloat should approxEqual(expectedLogProb) + it("logProb is -inf outside 0..K-1"): + val probs = Prob(Tensor(Shape(Axis[A] -> 4)).fromArray(Array(0.1f, 0.2f, 0.3f, 0.4f))) + + val dist = Categorical(probs) + // JAX would clamp 4 to the last index and wrap -1 around to it, giving log(0.4) + dist.logProb(Tensor0(4)).asFloat.item should be(Float.NegativeInfinity) + dist.logProb(Tensor0(-1)).asFloat.item should be(Float.NegativeInfinity) + it("sample distribution matches probabilities"): val probs = Prob(Tensor(Shape(Axis[A] -> 4)).fromArray(Array(0.1f, 0.2f, 0.3f, 0.4f))) val categorical = Categorical(probs) @@ -454,6 +514,15 @@ class DistributionSuite extends DimwitTest: ) scalaLogProbs.asFloat should approxEqual(jaxLogProbs) + it("logProb is -inf outside [0, 1]"): + val alpha = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(2.0f, 2.0f)) + val beta = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(2.0f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(-0.5f, 1.5f)) + + val dist = Beta(alpha, beta) + val logProbs = dist.elementWiseLogProb(x) + logProbs.asFloat.toArray.foreach(v => v should be(Float.NegativeInfinity)) + it("sample means approximates expected means"): val betaDist = Beta( Tensor(Shape(Axis[A] -> 2)).fromArray(Array(2.0f, 5.0f)), @@ -479,6 +548,14 @@ class DistributionSuite extends DimwitTest: ) scalaLogProbs.asFloat should approxEqual(jaxLogProbs) + it("logProb is -inf for x < 0"): + val rate = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.5f, 2.0f)) + val x = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(-1.0f, -0.1f)) + + val dist = Exponential(rate) + val logProbs = dist.elementWiseLogProb(x) + logProbs.asFloat.toArray.foreach(v => v should be(Float.NegativeInfinity)) + it("sample means approximates 1/rate"): val exponential = Exponential( Tensor(Shape(Axis[A] -> 2)).fromArray(Array(0.5f, 2.0f)) @@ -502,6 +579,14 @@ class DistributionSuite extends DimwitTest: ) scalaLogProbs.asFloat should approxEqual(jaxLogProbs) + it("logProb is -inf for x < 0"): + val rate = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 3.0f)) + val x = Tensor(Shape(Axis[A] -> 2)).fromArray(Array(-1, -2)) + + val dist = Poisson(rate) + val logProbs = dist.elementWiseLogProb(x) + logProbs.asFloat.toArray.foreach(v => v should be(Float.NegativeInfinity)) + it("sample means approximates rate"): val poisson = Poisson( Tensor(Shape(Axis[A] -> 2)).fromArray(Array(1.0f, 5.0f))