code: def test_mix_of_diag_normals_shared_cov_log_prob(batch_size):
      locs = torch.tensor([[-1.0, -1.0], [1.0, 1.0]])
      sigmas = torch.tensor([2.0, 2.0])
      logits = torch.tensor([math.log(0.25), math.log(0.75)])
      value = torch.tensor([0.5, 0.5])
      if batch_size > 1:
          locs = locs.unsqueeze(0).expand(batch_size, 2, 2)
          sigmas = sigmas.unsqueeze(0).expand(batch_size, 2)
          logits = logits.unsqueeze(0).expand(batch_size, 2)
          value = value.unsqueeze(0).expand(batch_size, 2)
      dist = MixtureOfDiagNormalsSharedCovariance(locs, sigmas, logits)
      log_prob = dist.log_prob(value)
      correct_log_prob = 0.25 * math.exp(-2.25 / 4.0)
      correct_log_prob += 0.75 * math.exp(-0.25 / 4.0)
      correct_log_prob /= 8.0 * math.pi
      correct_log_prob = math.log(correct_log_prob)
      if batch_size > 1:
          correct_log_prob = [correct_log_prob] * batch_size
      correct_log_prob = torch.tensor(correct_log_prob)
      assert_equal(
          log_prob,
          correct_log_prob,
          msg="bad log prob for MixtureOfDiagNormalsSharedCovariance",
      )
  
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 2 for calling a global function with example usage: test_mix_of_diag_normals_shared_cov_log_prob(batch_size) without return types