code: def expand_z_where(z_where):
      # Take a batch of three-vectors, and massages them into a batch of
      # 2x3 matrices with elements like so:
      # [s,x,y] -> [[s,0,x],
      #             [0,s,y]]
      n = z_where.size(0)
      out = torch.cat((z_where.new_zeros(n, 1), z_where), 1)
      ix = expansion_indices
      if z_where.is_cuda:
          ix = ix.cuda()
      out = torch.index_select(out, 1, ix)
      out = out.view(n, 2, 3)
      return out
  
imports: from collections import namedtuple
  import torch
  import torch.nn
  import torch.nn.functional
  from modules import MLP, Decoder, Encoder, Identity, Predict
  import pyro
  import pyro.distributions
label_desc: Use image node_id 2 for calling a global function with example usage: expand_z_where(z_where) and returns: out