add velocity and multi-factorization

Author: victor-shepardson

notepredictor/notepredictor/distributions.py

@@ -48,7 +48,12 @@ def get_params(self, h):
 
 
     def forward(self, h, x):
-        """log prob of x under distribution parameterized by h"""
+        """log prob of x under distribution parameterized by h
+        Args:
+            h: Tensor[...,n_params]
+            x: Tensor[...]
+            "..." dims must broadcast
+        """
         log_pi, loc, s = self.get_params(h)    
 
         d = self.res/2
@@ -59,8 +64,8 @@ def forward(self, h, x):
         # # censoring
         lo_cens = x <= self.lo+d
         hi_cens = x >= self.hi-d
-        ones = torch.ones_like(x_)
-        zeros = torch.zeros_like(x_)
+        ones = torch.ones_like(s)
+        zeros = torch.zeros_like(s)
 
         diff_term = torch.where(lo_cens | hi_cens, 
             ones, sd.exp() - (-sd).exp()

notepredictor/notepredictor/model.py

@@ -206,16 +206,20 @@ def mask_cat(missing, present, mask):
         # chunk into 7 and discard unmasked positions;
         # stack the masked positions along new first dim
         pitch_params, time_params, vel_params = (
-            torch.cat([
-                ch[None].expand(m, -1, -1, -1) 
+            torch.stack([
+                ch
                 for m,ch in zip(mask, dp.chunk(7, 0)) if m>0
                 ], 0)
             for mask,dp in zip(masks, dist_params)
         )
 
+        #TODO: weighting
+        # weights = np.log([[m for m in mask if m>0] for mask in masks]) # 3 x 4
+
         # get likelihoods
         pitch_logits = F.log_softmax(pitch_params, -1)
-        pitch_targets = pitches[None,:,1:,None] #1, batch, time-1, 1
+        # TODO: is gather working right with extra dim?
+        pitch_targets = pitches[None,:,1:,None].expand(4, -1, -1, -1) #1, batch, time-1, 1
         pitch_log_probs = pitch_logits.gather(-1, pitch_targets)[...,0]
 
         time_targets = times[:,1:]# batch, time-1