def test_mode_with_batch(self):
p = np.array([[0.2, 0.8], [0.6, 0.4]], dtype=np.float32)
logits = np.log(p) - 50.
dist = prob.Categorical(logits)
# test mode
modes = dist.mode()
self.assertArrayEqual([1, 0], modes)
def test_entropy_without_batch(self):
p = np.array([0.2, 0.8], dtype=np.float32)
logits = np.log(p) - 50.
dist = prob.Categorical(logits)
# test entropy
self.assertAllClose(-(0.2 * np.log(0.2) + 0.8 * np.log(0.8)),
dist.entropy(),
atol=0,
rtol=1e-5)
def make_categorical(batch_shape,
num_classes,
dtype=tf.int32,
seed=get_test_seed()):
logits = -50. + tf.random.uniform(list(batch_shape) + [num_classes],
minval=-10,
maxval=10,
dtype=tf.float32,
seed=seed)
return prob.Categorical(logits, dtype=dtype)
def test_entropy_with_batch(self):
p = np.array([[0.2, 0.8], [0.4, 0.6]], dtype=np.float32)
logits = np.log(p) - 50.
dist = prob.Categorical(logits)
self.assertAllClose([
-(0.2 * np.log(0.2) + 0.8 * np.log(0.8)),
-(0.4 * np.log(0.4) + 0.6 * np.log(0.6))
],
dist.entropy(),
atol=0,
rtol=1e-5)
def test_sample_with_batch(self):
p = np.array([[0.2, 0.8], [0.4, 0.6]], dtype=np.float32)
logits = np.log(p) - 50.
dist = prob.Categorical(logits)
# test sample
set_test_seed()
draws = np.asarray(dist.sample(10000))
self.assertFalse(np.any(draws < 0))
self.assertFalse(np.any(draws > 1))
self.assertAllClose(np.mean(draws == 0, axis=0), [0.2, 0.4], atol=1e-2)
self.assertAllClose(np.mean(draws == 1, axis=0), [0.8, 0.6], atol=1e-2)
def test_prob_with_batch(self):
p = np.array([[0.2, 0.8], [0.4, 0.6]], dtype=np.float32)
logits = np.log(p) - 50.
dist = prob.Categorical(logits)
# test logits
self.assertArrayEqual(dist.logits, logits)
self.assertArrayEqual(dist.logits.shape, p.shape)
# test _p, _log_p
self.assertAllClose(dist._p(), p)
# test prob, log_prob
self.assertAllClose(dist.prob([0, 0]), p[..., 0])
self.assertAllClose(dist.prob([1, 1]), p[..., 1])
def call(self, inputs, training=True):
'''Forward network
Args:
inputs (tf.Tensor): Expecting a latent vector in shape
(b, latent), tf.float32
training (bool, optional): Training mode. Defaults to True.
Returns:
Categorical: A categorical distribution
'''
return ub_prob.Categorical(self._model(inputs, training=training))
def test_entropy_with_neg_inf_logits(self):
probs = np.array([[0, 0.5, 0.5], [0, 1, 0]])
dist = prob.Categorical(np.log(probs))
ans = [-(0.5 * np.log(0.5) + 0.5 * np.log(0.5)), -(np.log(1))]
self.assertAllClose(ans, dist.entropy())
