Add random.binomial and random.multinomial

[carlosgmartin]

Add JAX counterparts of numpy.random.binomial and numpy.random.multinomial to jax.random package. See #480 (comment) for context. A current workaround is using the JAX substrate of TensorFlow Probability:

from jax import random, numpy as jnp
from tensorflow_probability.substrates import jax as tfp

def binomial(key, n, p, shape=()):
    return tfp.distributions.Binomial(n, probs=p).sample(
        seed=key,
        sample_shape=shape,
    )

def multinomial(key, n, p, shape=()):
    return tfp.distributions.Multinomial(n, probs=p).sample(
        seed=key,
        sample_shape=shape,
    )

key = random.PRNGKey(0)

key, subkey = random.split(key)
print(binomial(subkey, 9, .8, [2, 5]))

key, subkey = random.split(key)
print(multinomial(subkey, 9, jnp.array([.7, .2, .1]), [4]))

Output:

[[7. 8. 8. 7. 5.]
 [5. 8. 9. 8. 5.]]
[[7. 1. 1.]
 [4. 3. 2.]
 [5. 3. 1.]
 [4. 4. 1.]]

[zhangqiaorjc]

@sharadmv, is it possible to port tfp's implementation to jax.random?

[sharadmv]

Possibly, yes. The implementation is fairly complex though, and makes some accelerator-specific tradeoffs IIRC. cc: @srvasude @brianwa84.

[hylkedonker]

Any plans to add random.multinomial now that random.binomial has been merged?
The multinomial sampler could be implemented either as sequence of negative binomials or by repeated categorical draws.
See also the Wikipedia entry on multinomial variate generation).

[jakevdp]

Thanks for reaching out – I don't know of anyone working on this currently.

[andportnoy]

A workaround if you want pure JAX is to take the log of your probabilities vector (nonnegative, sums to 1):

jax.random.categorical(key, jnp.log(p))

(Based on the "contract" of jax.random.categorical:

logits: Unnormalized log probabilities of the categorical distribution(s) to sample from,
so that softmax(logits, axis) gives the corresponding probabilities.

)

That can't be the most efficient way to sample though...

[brianwa84]

Fyi: You would need to sum many categorical samples to get a multinomial sample.

…

On Sat, Jan 27, 2024, 7:23 PM Andrey Portnoy ***@***.***> wrote: A workaround if you want pure JAX is to take the log of your probabilities vector (that sums to 1): jax.random.categorical(key, jnp.log(p)) (Based on the "contract" of jax.random.categorical: logits: Unnormalized log probabilities of the categorical distribution(s) to sample from, so that softmax(logits, axis) gives the corresponding probabilities. ) That can't be the most efficient way to sample though... — Reply to this email directly, view it on GitHub <#13327 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AFJFSI62724M6LFSTMPHHGLYQWLABAVCNFSM6AAAAAASFQR3O6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSMJTGM4DGOJZGI> . You are receiving this because you were mentioned.Message ID: ***@***.***>