Normalizing distrax.Categorical

[patel-zeel]

Hi, I have noticed that distrax.Categorical normalizes the probabilities by their summation. For example,

import distrax
dist = distrax.Categorical(probs=[0.8, 1.2])
print(dist.probs)

output

[0.4 0.6]

This can go wrong at places, especially while fitting the parameters of this distribution to data. For example,

import distrax
dist = distrax.Categorical(probs=[-0.8, 1.2])
print(dist.probs)

output

[-1.9999999  2.9999998]

And now,

dist.sample(seed=jax.random.PRNGKey(0), sample_shape=(10,))

gives,

DeviceArray([-1, -1, -1, -1, -1, -1, -1, -1, -1, -1], dtype=int32)

I am not sure if this is done on purpose due to some benefits or this distribution can be improved to correct this behaviour.

[franrruiz]

Hi - This is intentional. The input probabilities are (of course) not allowed to be negative. To indicate that there's a mistake with the inputs, the samples are -1 when that is the case.

Note that the logits can be either negative or positive, so when fitting parameters to data, you may want to optimize with respect to the logits if you prefer to have unconstrained optimization.

I am closing the issue.

[patel-zeel]

Thank you for the clarification @franrruiz :)