def eager_binomial(total_count, probs, value):
probs = torch_stack((1 - probs, probs))
value = torch_stack((total_count - value, value))
return Multinomial(total_count, probs, value=value)
def test_torch_stack(n, shape, dim):
tensors = [torch.randn(shape) for _ in range(n)]
actual = torch_stack(tuple(Tensor(t) for t in tensors), dim=dim)
expected = Tensor(torch.stack(tensors, dim=dim))
assert_close(actual, expected)
def eager_beta(concentration1, concentration0, value):
concentration = torch_stack((concentration0, concentration1))
value = torch_stack((1 - value, value))
return Dirichlet(concentration, value=value)