original_docstring: 
          Encoder that ouputs parameters for base distribution of z and flow parameters.
          
docstring: Encoder that ouputs parameters for base distribution of z and flow parameters.
docstring_tokens: Encoder that ouputs parameters for base distribution of z and flow parameters .
code: def encode(self, x):
          """
          Encoder that ouputs parameters for base distribution of z and flow parameters.
          """
  
          batch_size = x.size(0)
  
          h = self.q_z_nn(x)
          h = h.view(-1, self.q_z_nn_output_dim)
          mean_z = self.q_z_mean(h)
          var_z = self.q_z_var(h)
  
          # Amortized r1, r2, q, b for all flows
  
          full_d = self.amor_d(h)
          diag1 = self.amor_diag1(h)
          diag2 = self.amor_diag2(h)
  
          full_d = full_d.resize(batch_size, self.num_ortho_vecs, self.num_ortho_vecs, self.num_flows)
          diag1 = diag1.resize(batch_size, self.num_ortho_vecs, self.num_flows)
          diag2 = diag2.resize(batch_size, self.num_ortho_vecs, self.num_flows)
  
          r1 = full_d * self.triu_mask
          r2 = full_d.transpose(2, 1) * self.triu_mask
  
          r1[:, self.diag_idx, self.diag_idx, :] = diag1
          r2[:, self.diag_idx, self.diag_idx, :] = diag2
  
          q = self.amor_q(h)
          b = self.amor_b(h)
  
          # Resize flow parameters to divide over K flows
          b = b.resize(batch_size, 1, self.num_ortho_vecs, self.num_flows)
  
          return mean_z, var_z, r1, r2, q, b
code_tokens: def encode ( self , x ) : batch_size = x . size ( 0 ) h = self . q_z_nn ( x ) h = h . view ( - 1 , self . q_z_nn_output_dim ) mean_z = self . q_z_mean ( h ) var_z = self . q_z_var ( h ) full_d = self . amor_d ( h ) diag1 = self . amor_diag1 ( h ) diag2 = self . amor_diag2 ( h ) full_d = full_d . resize ( batch_size , self . num_ortho_vecs , self . num_ortho_vecs , self . num_flows ) diag1 = diag1 . resize ( batch_size , self . num_ortho_vecs , self . num_flows ) diag2 = diag2 . resize ( batch_size , self . num_ortho_vecs , self . num_flows ) r1 = full_d * self . triu_mask r2 = full_d . transpose ( 2 , 1 ) * self . triu_mask r1 [ : , self . diag_idx , self . diag_idx , : ] = diag1 r2 [ : , self . diag_idx , self . diag_idx , : ] = diag2 q = self . amor_q ( h ) b = self . amor_b ( h ) b = b . resize ( batch_size , 1 , self . num_ortho_vecs , self . num_flows ) return mean_z , var_z , r1 , r2 , q , b
short_docstring: Encoder that ouputs parameters for base distribution of z and flow parameters.
short_docstring_tokens: Encoder that ouputs parameters for base distribution of z and flow parameters .
comment: """
          Encoder that ouputs parameters for base distribution of z and flow parameters.
          """$SEP$# Amortized r1, r2, q, b for all flows$SEP$# Resize flow parameters to divide over K flows