def METHOD_NAME(self):
    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)
    check((10,), (4,))
    check((2,3), (4,))
    check((3,4,5,6), (2,))