def test_categorical():
precision = 4
shape = (2, 3, 5)
ps = rng.random((np.prod(shape), 4))
ps = ps / np.sum(ps, axis=-1, keepdims=True)
data = np.reshape([rng.choice(4, p=p) for p in ps], shape)
ps = np.reshape(ps, shape + (4, ))
check_codec(shape, cs.Categorical(ps, precision), data)
def test_autoregressive():
precision = 8
batch_size = 3
data_size = 10
choices = 8
data = np.array(
[rng.choice(choices) for _ in range(batch_size * data_size)])
data = np.reshape(data, (batch_size, data_size))
fixed_probs = rng.random((batch_size, data_size, choices))
fixed_probs = fixed_probs / np.sum(fixed_probs, axis=-1, keepdims=True)
elem_idxs = [(slice(None), i)
for i in range(10)] # slice for the batch dimension
elem_codec = lambda p, idx: cs.Categorical(p, precision)
check_codec((batch_size, ),
cs.AutoRegressive(lambda *x: fixed_probs, (
batch_size,
data_size,
), fixed_probs.shape, elem_idxs, elem_codec), data)