def check(prob_shape, sample_shape):
for _ in range(4):
probs, probs2 = np.random.uniform(
0, 1, prob_shape), np.random.uniform(0, 1, prob_shape)
jc, jc2 = jd.Categorical(jt.array(probs)), jd.Categorical(
jt.array(probs2))
tc, tc2 = torch.distributions.Categorical(
torch.tensor(probs)), torch.distributions.Categorical(
torch.tensor(probs2))
assert np.allclose(
jc.entropy().data,
tc.entropy().numpy()), (jc.entropy().data,
tc.entropy().numpy())
x1 = jc.sample(sample_shape)
x2 = tc.sample(sample_shape)
assert tuple(x1.shape) == tuple(x2.shape)
x = np.random.randint(0, prob_shape[-1], tuple(x1.shape))
np.testing.assert_allclose(jc.log_prob(x),
tc.log_prob(torch.tensor(x)),
atol=1e-5)
np.testing.assert_allclose(jd.kl_divergence(jc, jc2),
torch.distributions.kl_divergence(
tc, tc2),
atol=1e-5)
def test_categorical(self):
import torch
for _ in range(10):
probs,probs2 = np.random.uniform(0,1,(10)), np.random.uniform(0,1,(10))
probs,probs2 = probs / probs.sum(),probs2 / probs2.sum()
jc, jc2 = jd.Categorical(jt.array(probs).reshape(1,-1)),jd.Categorical(jt.array(probs2).reshape(1,-1))
tc, tc2 = torch.distributions.Categorical(torch.tensor(probs)),torch.distributions.Categorical(torch.tensor(probs2))
assert np.allclose(jc.entropy().data,tc.entropy().numpy())
x = np.random.randint(0,10)
# print(jc.log_prob(x),tc.log_prob(x))
assert np.allclose(jc.log_prob(x),tc.log_prob(torch.tensor(x)))
assert np.allclose(jd.kl_divergence(jc,jc2),torch.distributions.kl_divergence(tc,tc2))
def test_cate(self):
a = jd.Categorical(jt.array([0.25, 0.25, 0.25, 0.25]))
x =np.array([0,0,0,0])
for i in range(1000):
x[a.sample().item()]+=1
assert (x > 200).all()
y = a.sample([2,3])
y.sync()
assert y.shape == [2,3]
def test_categorical1(self):
for _ in range(4):
probs, probs2 = np.random.uniform(0, 1, (10)), np.random.uniform(
0, 1, (10))
probs, probs2 = probs / probs.sum(), probs2 / probs2.sum()
jc, jc2 = jd.Categorical(jt.array(probs)), jd.Categorical(
jt.array(probs2))
tc, tc2 = torch.distributions.Categorical(
torch.tensor(probs)), torch.distributions.Categorical(
torch.tensor(probs2))
assert np.allclose(jc.entropy().data,
tc.entropy().numpy()), (jc.entropy().data,
tc.entropy().numpy())
x = np.random.randint(0, 10, (4))
np.testing.assert_allclose(jc.log_prob(x),
tc.log_prob(torch.tensor(x)),
atol=1e-5)
assert np.allclose(jd.kl_divergence(jc, jc2),
torch.distributions.kl_divergence(tc, tc2))