"logit" parameter of jax.random.categorical is misnamed, actually a log probability

summary

jax.random.categorical (introduced #1855) takes a parameter, logit, to specify the categorical distribution to sample from.

I believe this parameter is misnamed. It behaves like a log probability, not a logit.

potential resolutions

1. rename the parameter to logprob or something similar
2. change implementation and tests to behave like a logit
3. change to something less ambiguous, like a probability p

background

Maybe people use “logit” in different ways, but according to Wikipedia, logit(p) = ln( p / (1 - p) ). This is also the behavior implemented in scipy.special.logit (and jax.scipy.special.logit).

example / demo

import jax
import jax.numpy as np

key = jax.random.PRNGKey(seed=1)

p = np.array([0.4, 0.6])
n = 10000

arg = np.log(p / (1 - p))
# or:
# arg = jax.scipy.special.logit(p)
samples = jax.random.categorical(key, logits=arg, shape=(n,))

print(np.unique(samples, return_counts=True))
# (DeviceArray([0, 1], dtype=int32), DeviceArray([3059, 6941], dtype=int32))

arg = np.log(p)
samples = jax.random.categorical(key, logits=arg, shape=(n,))

print(np.unique(samples, return_counts=True))
# (DeviceArray([0, 1], dtype=int32), DeviceArray([3998, 6002], dtype=int32))

as you can see the counts for log(p) follow the given distribution, but the counts using logit(p) do not.

[shoyer]

If you like the conventions of TensorFlow probability, note that uou can already use TFP’s Categorical with JAX. We probably don’t advertise this well enough!

Alternatively, you can convert from probabilities to unnormalized log-probabilities just by using log(), e.g., jax.random.categorical(key, logits=jnp.log(p), shape=(n,)).

I would prefer to document either or both of these approaches in jax.random.categorical rather than to add a redundant argument.

[grisaitis]

thanks! yup, i was going to close.

thanks @shoyer about the fyi with tfp/jax. very useful!