code: def test_mean_gradient(
      K, D, flat_logits, cost_function, mix_dist, batch_mode
  ):
      n_samples = 200000
      if batch_mode:
          sample_shape = torch.Size(())
      else:
          sample_shape = torch.Size((n_samples,))
      if mix_dist == GaussianScaleMixture:
          locs = torch.zeros(K, D, requires_grad=True)
      else:
          locs = torch.rand(K, D).requires_grad_(True)
      if mix_dist == GaussianScaleMixture:
          component_scale = 1.5 * torch.ones(K) + 0.5 * torch.rand(K)
          component_scale.requires_grad_(True)
      else:
          component_scale = torch.ones(K, requires_grad=True)
      if mix_dist == MixtureOfDiagNormals:
          coord_scale = torch.ones(K, D) + 0.5 * torch.rand(K, D)
          coord_scale.requires_grad_(True)
      else:
          coord_scale = torch.ones(D) + 0.5 * torch.rand(D)
          coord_scale.requires_grad_(True)
      if not flat_logits:
          component_logits = (1.5 * torch.rand(K)).requires_grad_(True)
      else:
          component_logits = (0.1 * torch.rand(K)).requires_grad_(True)
      omega = (
          0.2 * torch.ones(D) + 0.1 * torch.rand(D)
      ).requires_grad_(False)
  
      _pis = torch.exp(component_logits)
      pis = _pis / _pis.sum()
  
      if cost_function == "cosine":
          analytic1 = torch.cos((omega * locs).sum(-1))
          analytic2 = torch.exp(
              -0.5
              * torch.pow(
                  omega * coord_scale * component_scale.unsqueeze(-1),
                  2.0,
              ).sum(-1)
          )
          analytic = (pis * analytic1 * analytic2).sum()
          analytic.backward()
      elif cost_function == "quadratic":
          analytic = torch.pow(
              coord_scale * component_scale.unsqueeze(-1), 2.0
          ).sum(-1) + torch.pow(locs, 2.0).sum(-1)
          analytic = (pis * analytic).sum()
          analytic.backward()
  
      analytic_grads = {}
      analytic_grads["locs"] = locs.grad.clone()
      analytic_grads["coord_scale"] = coord_scale.grad.clone()
      analytic_grads["component_logits"] = component_logits.grad.clone()
      analytic_grads["component_scale"] = component_scale.grad.clone()
  
      assert locs.grad.shape == locs.shape
      assert coord_scale.grad.shape == coord_scale.shape
      assert component_logits.grad.shape == component_logits.shape
      assert component_scale.grad.shape == component_scale.shape
  
      coord_scale.grad.zero_()
      component_logits.grad.zero_()
      locs.grad.zero_()
      component_scale.grad.zero_()
  
      if mix_dist == MixtureOfDiagNormalsSharedCovariance:
          params = {
              "locs": locs,
              "coord_scale": coord_scale,
              "component_logits": component_logits,
          }
          if batch_mode:
              locs = locs.unsqueeze(0).expand(n_samples, K, D)
              coord_scale = coord_scale.unsqueeze(0).expand(
                  n_samples, D
              )
              component_logits = component_logits.unsqueeze(0).expand(
                  n_samples, K
              )
              dist_params = {
                  "locs": locs,
                  "coord_scale": coord_scale,
                  "component_logits": component_logits,
              }
          else:
              dist_params = params
      elif mix_dist == MixtureOfDiagNormals:
          params = {
              "locs": locs,
              "coord_scale": coord_scale,
              "component_logits": component_logits,
          }
          if batch_mode:
              locs = locs.unsqueeze(0).expand(n_samples, K, D)
              coord_scale = coord_scale.unsqueeze(0).expand(
                  n_samples, K, D
              )
              component_logits = component_logits.unsqueeze(0).expand(
                  n_samples, K
              )
              dist_params = {
                  "locs": locs,
                  "coord_scale": coord_scale,
                  "component_logits": component_logits,
              }
          else:
              dist_params = params
      elif mix_dist == GaussianScaleMixture:
          params = {
              "coord_scale": coord_scale,
              "component_logits": component_logits,
              "component_scale": component_scale,
          }
          if batch_mode:
              return  # distribution does not support batched parameters
          else:
              dist_params = params
  
      dist = mix_dist(**dist_params)
      z = dist.rsample(sample_shape=sample_shape)
      assert z.shape == (n_samples, D)
      if cost_function == "cosine":
          cost = torch.cos((omega * z).sum(-1)).sum() / float(n_samples)
      elif cost_function == "quadratic":
          cost = torch.pow(z, 2.0).sum() / float(n_samples)
      cost.backward()
  
      assert_equal(
          analytic,
          cost,
          prec=0.1,
          msg="bad cost function evaluation for {} test (expected {}, got {})".format(
              mix_dist.__name__, analytic.item(), cost.item()
          ),
      )
      logger.debug(
          "analytic_grads_logit: {}".format(
              analytic_grads["component_logits"].detach().cpu().numpy()
          )
      )
  
      for param_name, param in params.items():
          assert_equal(
              param.grad,
              analytic_grads[param_name],
              prec=0.1,
              msg="bad {} grad for {} (expected {}, got {})".format(
                  param_name,
                  mix_dist.__name__,
                  analytic_grads[param_name],
                  param.grad,
              ),
          )
  
file_path: pyro/tests/distributions/test_gaussian_mixtures.py
imports: import logging
  import math
  import pytest
  import torch
  from pyro.distributions import GaussianScaleMixture, MixtureOfDiagNormals, MixtureOfDiagNormalsSharedCovariance
  from tests.common import assert_equal
label_desc: Use image node_id 1 for calling a global function with example usage: test_mean_gradient(K, D, flat_logits, cost_function, mix_dist, batch_mode) without return types