code: def test_rsample(dist):
      if not dist.pyro_dist.has_rsample:
          return
      for idx in range(len(dist.dist_params)):
          # Compute CPU value.
          with torch.device("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 torch.device("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())
  
imports: import pytest
  import torch
  from torch.autograd import grad
  from tests.common import assert_equal, requires_cuda, xfail_if_not_implemented
label_desc: Use image node_id 2 for calling a global function with example usage: test_rsample(dist) without return types