Skip to content

Return -inf from Categorical.logProb outside 0..K-1 #180

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
marcelluethi merged 1 commit into from
Oct 27, 2026
Merged

Return -inf from Categorical.logProb outside 0..K-1 #180

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion core/src/main/scala/dimwit/stats/UnivariateDistributions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)))

@benikm91 benikm91 Oct 26, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this should work as which read better (into cast):

val inRange = (x >= 0) and (x < 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))
Expand Down
85 changes: 85 additions & 0 deletions core/src/test/scala/dimwit/stats/DistributionSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -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)),
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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)),
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)),
Expand All @@ -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))
Expand All @@ -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))
Expand Down
Loading