code: def test_onehot_svi_usage():
      def model():
          p = torch.tensor([0.25] * 4)
          pyro.sample("z", OneHotCategorical(probs=p))
  
      def guide():
          q = pyro.param(
              "q",
              torch.tensor([0.1, 0.2, 0.3, 0.4]),
              constraint=constraints.simplex,
          )
          temp = torch.tensor(0.10)
          pyro.sample(
              "z",
              RelaxedOneHotCategoricalStraightThrough(
                  temperature=temp, probs=q
              ),
          )
  
      adam = optim.Adam({"lr": 0.001, "betas": (0.95, 0.999)})
      svi = SVI(model, guide, adam, loss=Trace_ELBO())
  
      for k in range(6000):
          svi.step()
  
      assert_equal(
          pyro.param("q"),
          torch.tensor([0.25] * 4),
          prec=0.01,
          msg="test svi usage of RelaxedOneHotCategoricalStraightThrough failed",
      )
  
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 3 for calling a global function with example usage: test_onehot_svi_usage() without return types