def _test(shape, n): K = shape[-1] rv = Multinomial(shape, pi=tf.constant(1.0/K, shape=shape)) rv_sample = rv.sample(n) x = rv_sample.eval() x_tf = tf.constant(x, dtype=tf.float32) pi = rv.pi.eval() if len(shape) == 1: assert np.allclose( rv.log_prob_idx((), x_tf).eval(), multinomial_logpmf_vec(x[:, :], 1, pi[:])) elif len(shape) == 2: for i in range(shape[0]): assert np.allclose( rv.log_prob_idx((i, ), x_tf).eval(), multinomial_logpmf_vec(x[:, i, :], 1, pi[i, :])) else: assert False
def _test(shape, p, size): x = Multinomial(shape, p) val_est = tuple(get_dims(x.sample(size=size))) val_true = (size, ) + shape assert val_est == val_true
def _test(shape, p, n): x = Multinomial(shape, p) val_est = tuple(get_dims(x.sample(n))) val_true = (n, ) + shape assert val_est == val_true