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