code: def test_shape_augmented_beta(alpha, beta):
      num_samples = 10000
      alphas = torch.full((num_samples, 1), alpha).requires_grad_()
      betas = torch.full((num_samples, 1), beta).requires_grad_()
      dist = ShapeAugmentedBeta(
          alphas, betas
      )  # implemented using Rejector
      z = dist.rsample()
      cost = z.sum()
      (cost + cost.detach() * dist.score_parts(z)[1]).backward()
      mean_alpha_grad = alphas.grad.mean().item()
      mean_beta_grad = betas.grad.mean().item()
      expected_alpha_grad = beta / (alpha + beta) ** 2
      expected_beta_grad = -alpha / (alpha + beta) ** 2
      assert_equal(
          mean_alpha_grad,
          expected_alpha_grad,
          prec=0.02,
          msg="bad grad for alpha",
      )
      assert_equal(
          mean_beta_grad,
          expected_beta_grad,
          prec=0.02,
          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 9 for calling a global function with example usage: test_shape_augmented_beta(alpha, beta) without return types