def test_simple_det(self):
for n_bins, n_samples, batch in product([7, 20], [2, 7, 31, 32, 33],
[(), (1, 4), (31, ), (32, ),
(33, )]):
weights = torch.rand(size=(batch + (n_bins, )))
bins = torch.cumsum(torch.rand(size=(batch + (n_bins + 1, ))),
dim=-1)
python = sample_pdf_python(bins, weights, n_samples, det=True)
cpp = sample_pdf(bins, weights, n_samples, det=True)
self.assertClose(cpp, python, atol=2e-3)
nthreads = torch.get_num_threads()
torch.set_num_threads(1)
cpp_singlethread = sample_pdf(bins, weights, n_samples, det=True)
self.assertClose(cpp_singlethread, python, atol=2e-3)
torch.set_num_threads(nthreads)
device = torch.device("cuda:0")
cuda = sample_pdf(bins.to(device),
weights.to(device),
n_samples,
det=True).cpu()
self.assertClose(cuda, python, atol=2e-3)
def test_rand_cpu(self):
n_bins, n_samples, batch_size = 11, 17, 9
weights = torch.rand(size=(batch_size, n_bins))
bins = torch.cumsum(torch.rand(size=(batch_size, n_bins + 1)), dim=-1)
torch.manual_seed(1)
python = sample_pdf_python(bins, weights, n_samples)
torch.manual_seed(1)
cpp = sample_pdf(bins, weights, n_samples)
self.assertClose(cpp, python, atol=2e-3)
def test_rand_nogap(self):
# Case where random is actually deterministic
weights = torch.FloatTensor([0, 10, 0])
bins = torch.FloatTensor([0, 10, 10, 25])
n_samples = 8
predicted = torch.full((n_samples, ), 10.0)
python = sample_pdf_python(bins, weights, n_samples)
self.assertClose(python, predicted)
cpp = sample_pdf(bins, weights, n_samples)
self.assertClose(cpp, predicted)
device = torch.device("cuda:0")
cuda = sample_pdf(bins.to(device), weights.to(device), n_samples).cpu()
self.assertClose(cuda, predicted)
def test_single_bin(self):
bins = torch.arange(2).expand(5, 2) + 17
weights = torch.ones(5, 1)
output = sample_pdf_python(bins, weights, 100, True)
calc = torch.linspace(17, 18, 100).expand(5, -1)
self.assertClose(output, calc)