class_docstr: Convolutional Layers to estimate NMF Activations from Classifier Representations
  
  Arguments
  ---------
  n_comp : int
      Number of NMF components (or equivalently number of neurons at the output per timestep)
  T: int
      The targeted length along the time dimension
  in_emb_dims: List with int elements
      A list with length 3 that contains the dimensionality of the input dimensions
      The list needs to match the number of channels in the input classifier representations
      The last entry should be the smallest entry
  
  Example
  -------
  >>> inp = [torch.ones(2, 150, 6, 2), torch.ones(2, 100, 6, 2), torch.ones(2, 50, 12, 5)]
  >>> psi = Psi(n_comp=100, T=120, in_emb_dims=[150, 100, 50])
  >>> h = psi(inp)
  >>> print(h.shape)
  torch.Size([2, 100, 120])
class_docstr_tok: ["Convolutional","Layers","to","estimate","NMF","Activations","from","Classifier","Representations","Arguments","--","--","--","--","-","n_comp",":","int","Number","of","NMF","components","(","or","equivalently","number","of","neurons","at","the","output","per","timestep",")","T",":","int","The","targeted","length","along","the","time","dimension","in_emb_dims",":","List","with","int","elements","A","list","with","length","3","that","contains","the","dimensionality","of","the","input","dimensions","The","list","needs","to","match","the","number","of","channels","in","the","input","classifier","representations","The","last","entry","should","be","the","smallest","entry","Example","--","--","--","-",">",">",">","inp","=","[","torch.ones","(","2",",","150",",","6",",","2",")",",","torch.ones","(","2",",","100",",","6",",","2",")",",","torch.ones","(","2",",","50",",","12",",","5",")","]",">",">",">","psi","=","Psi","(","n_comp=100",",","T=120",",","in_emb_dims=","[","150",",","100",",","50","]",")",">",">",">","h","=","psi","(","inp",")",">",">",">","print","(","h.shape",")","torch.Size","(","[","2",",","100",",","120","]",")"]
docstr: This forward function returns the NMF time activations given classifier activations
  Arguments
  ---------
      inp: A length 3 list of classifier input representions.
docstr_tok: ["This","forward","function","returns","the","NMF","time","activations","given","classifier","activations","Arguments","--","--","--","--","-","inp",":","A","length","3","list","of","classifier","input","representions","."]
code: def forward(self, inp):
      """This forward function returns the NMF time activations given classifier activations
      Arguments
      ---------
          inp: A length 3 list of classifier input representions.
      """
      error = "in PSI doesn't match. The embedding dimensions need to be consistent with the list self.in_emb_dims"
      for i, in_emb_dim in enumerate(self.in_emb_dims):
          # sanity check on shapes
          assert inp[i].shape[1] == self.in_emb_dims[i], (
              "Nr. of channels " + error
          )
  
      assert inp[0].shape[2] == inp[1].shape[2], (
          "Spatial dimension " + error
      )
      assert inp[0].shape[3] == inp[1].shape[3], (
          "Spatial dimension " + error
      )
      assert 2 * inp[0].shape[3] == (inp[2].shape[3] - 1), (
          "Spatial dimension "
          + error
          + f" 1st (idx 0) element has shape {inp[0].shape[3]} second element (idx 1) has shape {inp[2].shape[3]}"
      )
  
      x1, x2, x3 = inp
  
      # upsample inp[0] and inp[1] time and frequency axis once
      x1 = self.upsamp(x1)
      x2 = self.upsamp(x2)
  
      # compress feature number to the min among given hidden representations
      x1 = self.act(self.c1(x1))
      x2 = self.act(self.c2(x2))
  
      # for compatibility with cnn14 fixed frequency dimension
      x1 = F.pad(x1, (0, 1, 0, 0))
      x2 = F.pad(x2, (0, 1, 0, 0))
      x = torch.cat((x1, x2, x3), axis=1)
  
      # upsample time axis and collapse freq
      x = self.upsamp_time(x)
  
      # mix contribution for the three hidden layers -- work on this when fixing training
      x = self.conv(x)
      x = self.act(self.out_conv(x)).squeeze(3)
      return x
  
imports: import torch.nn
  import torch.nn.functional
  import torch
  from speechbrain.lobes.models.PIQ import ResBlockAudio
label_desc: Use image node_id 2 for calling the Psi obj's underlying member method code with example usage: obj.forward(inp) and returns: x