def sample(self, key, size=()):
return categorical(key, self.probs, shape=size + self.batch_shape)
def test_categorical_shape(p, shape):
rng_key = random.PRNGKey(0)
expected_shape = lax.broadcast_shapes(p.shape[:-1], shape)
assert jnp.shape(categorical(rng_key, p, shape)) == expected_shape
def test_categorical_stats(p):
rng_key = random.PRNGKey(0)
n = 10000
z = categorical(rng_key, p, (n,))
_, counts = np.unique(z, return_counts=True)
assert_allclose(counts / float(n), p, atol=0.01)
def sample(self, key, sample_shape=()):
assert is_prng_key(key)
return categorical(key,
self.probs,
shape=sample_shape + self.batch_shape)
