code: def test_bernoulli_entropy_grad(temp):
      num_samples = 1500000
      q = torch.tensor([0.1, 0.2, 0.3, 0.4], requires_grad=True)
      temp = torch.tensor(temp)
  
      dist_q = RelaxedBernoulli(temperature=temp, probs=q)
      z = dist_q.rsample(sample_shape=(num_samples,))
      expected = grad(dist_q.log_prob(z).sum(), [q])[0] / num_samples
  
      dist_q = RelaxedBernoulliStraightThrough(
          temperature=temp, probs=q
      )
      z = dist_q.rsample(sample_shape=(num_samples,))
      actual = grad(dist_q.log_prob(z).sum(), [q])[0] / num_samples
  
      assert_equal(
          expected,
          actual,
          prec=0.04,
          msg="bad grad for RelaxedBernoulliStraightThrough (expected {}, got {})".format(
              expected, actual
          ),
      )
  
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 5 for calling a global function with example usage: test_bernoulli_entropy_grad(temp) without return types