code: def test_shape_augmented_gamma_elbo(alpha, beta):
      num_samples = 100000
      alphas = torch.full((num_samples, 1), alpha).requires_grad_()
      betas = torch.full((num_samples, 1), beta).requires_grad_()
  
      model = Gamma(
          torch.ones(num_samples, 1), torch.ones(num_samples, 1)
      )
      guide1 = Gamma(alphas, betas)
      guide2 = ShapeAugmentedGamma(
          alphas, betas
      )  # implemented using Rejector
  
      grads = []
      for guide in [guide1, guide2]:
          grads.append(compute_elbo_grad(model, guide, [alphas, betas]))
      expected, actual = grads
      expected = [g.mean() for g in expected]
      actual = [g.mean() for g in actual]
      scale = [(1 + abs(g)) for g in expected]
      assert_equal(
          actual[0] / scale[0],
          expected[0] / scale[0],
          prec=0.05,
          msg="bad grad for alpha",
      )
      assert_equal(
          actual[1] / scale[1],
          expected[1] / scale[1],
          prec=0.05,
          msg="bad grad for beta",
      )
  
file_path: pyro/tests/distributions/test_rejector.py
imports: import pytest
  import torch
  from torch.autograd import grad
  from pyro.distributions import Exponential, Gamma
  from pyro.distributions.testing.rejection_exponential import RejectionExponential
  from pyro.distributions.testing.rejection_gamma import RejectionGamma, RejectionStandardGamma, ShapeAugmentedBeta, ShapeAugmentedGamma
  from tests.common import assert_equal
label_desc: Use image node_id 8 for calling a global function with example usage: test_shape_augmented_gamma_elbo(alpha, beta) without return types