code: def test_bernoulli_shapes(probs):
      temperature = torch.tensor(0.5)
      probs = torch.tensor(probs, requires_grad=True)
      d = RelaxedBernoulliStraightThrough(temperature, probs=probs)
      sample = d.rsample()
      log_prob = d.log_prob(sample)
      grad_probs = grad(log_prob.sum(), [probs])[0]
      assert grad_probs.shape == probs.shape
  
file_path: pyro/tests/distributions/test_relaxed_straight_through.py
imports: import pytest
  import torch
  from torch.autograd import grad
  from torch.distributions import constraints
  import pyro
  import pyro.optim
  from pyro.distributions import OneHotCategorical, RelaxedBernoulli, RelaxedBernoulliStraightThrough, RelaxedOneHotCategorical, RelaxedOneHotCategoricalStraightThrough
  from pyro.infer import SVI, Trace_ELBO
  from tests.common import assert_equal
label_desc: Use image node_id 4 for calling a global function with example usage: test_bernoulli_shapes(probs) without return types