def test_log_prob(dist): for idx in range(len(dist.dist_params)): # Compute CPU value. with tensors_default_to("cpu"): data = dist.get_test_data(idx) params = dist.get_dist_params(idx) with xfail_if_not_implemented(): cpu_value = dist.pyro_dist(**params).log_prob(data) assert not cpu_value.is_cuda # Compute GPU value. with tensors_default_to("cuda"): data = dist.get_test_data(idx) params = dist.get_dist_params(idx) cuda_value = dist.pyro_dist(**params).log_prob(data) assert cuda_value.is_cuda assert_equal(cpu_value, cuda_value.cpu())
def test_log_prob(dist): for idx in range(len(dist.dist_params)): # Compute CPU value. with tensors_default_to("cpu"): data = dist.get_test_data(idx) params = dist.get_dist_params(idx) with xfail_if_not_implemented(): cpu_value = dist.pyro_dist(**params).log_prob(data) assert not cpu_value.is_cuda # Compute GPU value. with tensors_default_to("cuda"): data = dist.get_test_data(idx) params = dist.get_dist_params(idx) cuda_value = dist.pyro_dist(**params).log_prob(data) assert cuda_value.is_cuda assert_equal(cpu_value, cuda_value.cpu())
def test_sample(dist): for idx in range(len(dist.dist_params)): # Compute CPU value. with tensors_default_to("cpu"): params = dist.get_dist_params(idx) try: cpu_value = dist.pyro_dist.sample(**params) except ValueError as e: pytest.xfail('CPU version fails: {}'.format(e)) assert not cpu_value.is_cuda # Compute GPU value. with tensors_default_to("cuda"): params = dist.get_dist_params(idx) cuda_value = dist.pyro_dist.sample(**params) assert cuda_value.is_cuda assert_equal(cpu_value.size(), cuda_value.size())
def test_rsample(dist): if not dist.pyro_dist.has_rsample: return for idx in range(len(dist.dist_params)): # Compute CPU value. with tensors_default_to("cpu"): params = dist.get_dist_params(idx) grad_params = [ key for key, val in params.items() if torch.is_tensor(val) and val.dtype in (torch.float32, torch.float64) ] for key in grad_params: val = params[key].clone() val.requires_grad = True params[key] = val try: with xfail_if_not_implemented(): cpu_value = dist.pyro_dist(**params).rsample() cpu_grads = grad(cpu_value.sum(), [params[key] for key in grad_params]) except ValueError as e: pytest.xfail('CPU version fails: {}'.format(e)) assert not cpu_value.is_cuda # Compute GPU value. with tensors_default_to("cuda"): params = dist.get_dist_params(idx) for key in grad_params: val = params[key].clone() val.requires_grad = True params[key] = val cuda_value = dist.pyro_dist(**params).rsample() assert cuda_value.is_cuda assert_equal(cpu_value.size(), cuda_value.size()) cuda_grads = grad(cuda_value.sum(), [params[key] for key in grad_params]) for cpu_grad, cuda_grad in zip(cpu_grads, cuda_grads): assert_equal(cpu_grad.size(), cuda_grad.size())
def test_rsample(dist): if not dist.pyro_dist.has_rsample: return for idx in range(len(dist.dist_params)): # Compute CPU value. with tensors_default_to("cpu"): params = dist.get_dist_params(idx) grad_params = [key for key, val in params.items() if torch.is_tensor(val) and val.dtype in (torch.float32, torch.float64)] for key in grad_params: val = params[key].clone() val.requires_grad = True params[key] = val try: with xfail_if_not_implemented(): cpu_value = dist.pyro_dist(**params).rsample() cpu_grads = grad(cpu_value.sum(), [params[key] for key in grad_params]) except ValueError as e: pytest.xfail('CPU version fails: {}'.format(e)) assert not cpu_value.is_cuda # Compute GPU value. with tensors_default_to("cuda"): params = dist.get_dist_params(idx) for key in grad_params: val = params[key].clone() val.requires_grad = True params[key] = val cuda_value = dist.pyro_dist(**params).rsample() assert cuda_value.is_cuda assert_equal(cpu_value.size(), cuda_value.size()) cuda_grads = grad(cuda_value.sum(), [params[key] for key in grad_params]) for cpu_grad, cuda_grad in zip(cpu_grads, cuda_grads): assert_equal(cpu_grad.size(), cuda_grad.size())