code: def backward(ctx, grad_output):
      z, coord_scale, locs, component_logits, pis = ctx.saved_tensors
      K = component_logits.size(-1)
      batch_dims = coord_scale.dim() - 1
      g = grad_output  # l b i
  
      z_tilde = z / coord_scale  # l b i
      locs_tilde = locs / coord_scale.unsqueeze(-2)  # b j i
      mu_ab = locs_tilde.unsqueeze(-2) - locs_tilde.unsqueeze(
          -3
      )  # b k j i
      mu_ab_norm = torch.pow(mu_ab, 2.0).sum(-1).sqrt()  # b k j
      mu_ab /= mu_ab_norm.unsqueeze(-1)  # b k j i
      diagonals = torch.empty((K,), dtype=torch.long, device=z.device)
      torch.arange(K, out=diagonals)
      mu_ab[..., diagonals, diagonals, :] = 0.0
  
      mu_ll_ab = (locs_tilde.unsqueeze(-2) * mu_ab).sum(-1)  # b k j
      z_ll_ab = (z_tilde.unsqueeze(-2).unsqueeze(-2) * mu_ab).sum(
          -1
      )  # l b k j
      z_perp_ab = (
          z_tilde.unsqueeze(-2).unsqueeze(-2)
          - z_ll_ab.unsqueeze(-1) * mu_ab
      )  # l b k j i
      z_perp_ab_sqr = torch.pow(z_perp_ab, 2.0).sum(-1)  # l b k j
  
      epsilons = z_tilde.unsqueeze(-2) - locs_tilde  # l b j i
      log_qs = -0.5 * torch.pow(epsilons, 2.0)  # l b j i
      log_q_j = log_qs.sum(-1, keepdim=True)  # l b j 1
      log_q_j_max = torch.max(log_q_j, -2, keepdim=True)[0]
      q_j_prime = torch.exp(log_q_j - log_q_j_max)  # l b j 1
      q_j = torch.exp(log_q_j)  # l b j 1
  
      q_tot = (pis.unsqueeze(-1) * q_j).sum(-2)  # l b 1
      q_tot_prime = (
          (pis.unsqueeze(-1) * q_j_prime).sum(-2).unsqueeze(-1)
      )  # l b 1 1
  
      root_two = math.sqrt(2.0)
      mu_ll_ba = torch.transpose(mu_ll_ab, -1, -2)
      logits_grad = torch.erf(
          (z_ll_ab - mu_ll_ab) / root_two
      ) - torch.erf((z_ll_ab + mu_ll_ba) / root_two)
      logits_grad *= torch.exp(-0.5 * z_perp_ab_sqr)  # l b k j
  
      #                 bi      lbi                               bkji
      mu_ab_sigma_g = (
          (coord_scale * g).unsqueeze(-2).unsqueeze(-2) * mu_ab
      ).sum(
          -1
      )  # l b k j
      logits_grad *= -mu_ab_sigma_g * pis.unsqueeze(-2)  # l b k j
      logits_grad = pis * sum_leftmost(
          logits_grad.sum(-1) / q_tot, -(1 + batch_dims)
      )  # b k
      logits_grad *= math.sqrt(0.5 * math.pi)
  
      #           b j                 l b j 1   l b i             l b 1 1
      prefactor = (
          pis.unsqueeze(-1) * q_j_prime * g.unsqueeze(-2) / q_tot_prime
      )  # l b j i
      locs_grad = sum_leftmost(prefactor, -(2 + batch_dims))  # b j i
      coord_scale_grad = sum_leftmost(
          prefactor * epsilons, -(2 + batch_dims)
      ).sum(
          -2
      )  # b i
  
      return locs_grad, coord_scale_grad, logits_grad, None, None, None
  
file_path: pyro/pyro/distributions/diag_normal_mixture_shared_cov.py
imports: import math
  import torch
  from torch.autograd import Function
  from torch.autograd.function import once_differentiable
  from torch.distributions import Categorical, constraints
  from pyro.distributions.torch_distribution import TorchDistribution
  from pyro.distributions.util import sum_leftmost
label_desc: Use image node_id 2 for calling the _MixDiagNormalSharedCovarianceSample obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: locs_grad,  coord_scale_grad,  logits_grad,  None,  None,  None