def test_one_hot(self):
a = jd.OneHotCategorical(jt.array([0.25, 0.25, 0.25, 0.25]))
x = a.sample().numpy()
for i in range(1000):
x += a.sample().numpy()
assert (x > 200).all()
y = a.sample([2, 3])
y.sync()
assert y.shape == [2, 3, 4]
probs, probs2 = np.random.uniform(0, 1,
(10)), np.random.uniform(0, 1, (10))
probs, probs2 = probs / probs.sum(), probs2 / probs2.sum()
import torch
jc, jc2 = jd.OneHotCategorical(jt.array(probs).reshape(
1, -1)), jd.OneHotCategorical(jt.array(probs2).reshape(1, -1))
tc, tc2 = torch.distributions.OneHotCategorical(
torch.tensor(probs)), torch.distributions.OneHotCategorical(
torch.tensor(probs2))
assert np.allclose(jc.entropy().data, tc.entropy().numpy())
x = np.zeros((4, 10))
for _ in range(4):
nx = np.random.randint(0, 9)
x[_, nx] = 1
assert np.allclose(jc.log_prob(jt.array(x)),
tc.log_prob(torch.tensor(x)))
assert np.allclose(jd.kl_divergence(jc, jc2),
torch.distributions.kl_divergence(tc, tc2))
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.OneHotCategorical(
jt.array(probs)), jd.OneHotCategorical(jt.array(probs2))
tc, tc2 = torch.distributions.OneHotCategorical(
torch.tensor(
probs)), torch.distributions.OneHotCategorical(
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_geometric(self):
import torch
for _ in range(4):
prob, prob2 = np.random.uniform(0, 1), np.random.uniform(0, 1)
jg, jg2 = jd.Geometric(prob), jd.Geometric(prob2)
tg, tg2 = torch.distributions.Geometric(
prob), torch.distributions.Geometric(prob2)
assert np.allclose(jg.entropy().data, tg.entropy().numpy())
x = np.random.randint(1, 10)
assert np.allclose(jg.log_prob(x), tg.log_prob(torch.tensor(x)))
# print(jd.kl_divergence(jg,jg2),torch.distributions.kl_divergence(tg,tg2))
assert np.allclose(jd.kl_divergence(jg, jg2),
torch.distributions.kl_divergence(tg, tg2))
def test_uniform(self):
for _ in range(4):
low, low2 = np.random.randint(-1, 2), np.random.randint(-1, 2)
leng, leng2 = np.random.uniform(0, 2), np.random.uniform(0, 2)
high, high2 = low + leng, low2 + leng2
ju, ju2 = jd.Uniform(low, high), jd.Uniform(low2, high2)
tu, tu2 = torch.distributions.Uniform(
low, high), torch.distributions.Uniform(low2, high2)
assert np.allclose(ju.entropy().data, tu.entropy().numpy())
x = np.random.uniform(low, high)
assert np.allclose(ju.log_prob(x), tu.log_prob(torch.tensor(x)))
assert np.allclose(jd.kl_divergence(ju, ju2),
torch.distributions.kl_divergence(tu, tu2))
def test_normal(self):
for _ in range(4):
mu = np.random.uniform(-1, 1)
sigma = np.random.uniform(0, 2)
jn = jd.Normal(mu, sigma)
tn = torch.distributions.Normal(mu, sigma)
assert np.allclose(jn.entropy().data, tn.entropy().numpy())
x = np.random.uniform(-1, 1)
assert np.allclose(jn.log_prob(x), tn.log_prob(torch.tensor(x)))
mu2 = np.random.uniform(-1, 1)
sigma2 = np.random.uniform(0, 2)
jn2 = jd.Normal(mu2, sigma2)
tn2 = torch.distributions.Normal(mu2, sigma2)
assert np.allclose(
jd.kl_divergence(jn, jn2).data,
torch.distributions.kl_divergence(tn, tn2).numpy())
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))
