Repository navigation
Return -inf from Categorical.logProb outside 0..K-1 - #180
Conversation
JAX does not raise on out-of-range indices: it clamps an index >= K to K-1 and wraps a negative index around. So Categorical.logProb returned a finite, plausible log-probability for impossible values, e.g. log p(K-1) for x = K or x = -1. Values outside 0..K-1 are now masked to -inf. Add tests ensuring that logprob is -inf outside its support
| 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))) |
There was a problem hiding this comment.
this should work as which read better (into cast):
val inRange = (x >= 0) and (x < probs.shape.dimensions(0))
There was a problem hiding this comment.
I like the change. It is in DimWit's spirit to not just wrap values around.
I am fine with "-inf" but did you consider to throw a runtime error? Maybe not possible because we are in tensorland?
The code can be slightly improved but good as it is. Approved.
It is a good point and we should maybe think a bit more deeply about it. My hunch is that |
Return -inf from Categorical.logProb outside 0..K-1
JAX does not raise on out-of-range indices: it clamps an index >= K to K-1 and wraps a negative index around.
So Categorical.logProb returned a finite, plausible log-probability for impossible values, e.g. log p(K-1)for x = K or x = -1. Values outside 0..K-1 are now masked to -inf.
It should be a law in dimwit that logprob of any value outside the support is
-inf(log(0))Tests are added, which ensure that logprob is -inf outside its support