jax-ml / jax-ml/oryx

expit of logistic variable gives log_prob 0.0

Open
#65 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
328
Forks
13
Avg merge
49m
Merged PRs (30d)
2

Description

I'm trying to get the logprob of a expit-transformed logistic-distributed variable but it always returns zero:

```
import jax.random
from oryx.core.ppl import random_variable, log_prob
from jax.scipy.special import expit
import oryx.distributions as tfd

def simple_sample(key):
a = random_variable(tfd.Logistic(0., 1.))(key)
return expit(a)

x = simple_sample(jax.random.PRNGKey(0))
print(x) # 0.41845703
print(log_prob(simple_sample)(0.5)) # 0.0
print(log_prob(simple_sample)(x)) # 0.0
```
Versions:
jax-0.4.25
oryx-0.2.6

(Both a exp transformed logistic variable and a expit transformed normal variable seems to work, so there is something special about this combination)

Contributor guide

Open the contributing guide

Research direction

Run the provided Python reproducer with jax, oryx.core.ppl.random_variable and log_prob, then compare it with the working exp-transformed logistic and expit-transformed normal cases mentioned in the report. Trace how Logistic and expit are handled by log_prob; done means the expit-transformed logistic variable produces the correct log probability for both 0.5 and the sampled value.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.