code: def forward(
      ctx, locs, coord_scale, component_logits, pis, which, noise_shape
  ):
      dim = coord_scale.size(-1)
      white = locs.new_empty(noise_shape).normal_()
      n_unsqueezes = locs.dim() - which.dim()
      for _ in range(n_unsqueezes):
          which = which.unsqueeze(-1)
      expand_tuple = tuple(which.shape[:-1] + (dim,))
      loc = torch.gather(locs, -2, which.expand(expand_tuple)).squeeze(
          -2
      )
      z = loc + coord_scale * white
      ctx.save_for_backward(z, coord_scale, locs, component_logits, pis)
      return z
  
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 1 for calling the _MixDiagNormalSharedCovarianceSample obj's underlying member method code with example usage: obj.forward(ctx, locs, coord_scale, component_logits, pis, which, noise_shape) and returns: z