Hi everybody,
I was wondering if there would be a way in PyTorch to compute dynamic causal convolution: where the kernel size is different per sample?
Let’s take the simple case where the kernel_size is the same for every sample in the batch:
In [1]: import torch       
   ...: import torch.nn as nn 
   ...: import torch.nn.functional as F                                                                                                                                                                                                                                                                                                                                                                                                                                                                            

In [2]: inputs1 = [0 for _ in range(20)]     
   ...: inputs1[8] = 1 
   ...: inputs2 = [0 for _ in range(20)] 
   ...: inputs2[13] = 1 
   ...: inputs = torch.FloatTensor([inputs1, inputs2])                                                                                                                                                                                                                           

In [3]: kernel_size = 6 # Known a priori 
   ...: kernels = torch.cat([torch.ones(kernel_size), torch.zeros(kernel_size - 1)])                                                                                                                                                                                             

In [4]: inputs = inputs.unsqueeze(1) 
   ...: inputs                                                                                                                                                                                                                                                                   
Out[4]: 
tensor([[[0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.,
          0., 0., 0.]],

        [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,
          0., 0., 0.]]])

In [5]: kernels = kernels.unsqueeze(0).unsqueeze(0) 
   ...: kernels # torch.Size([1, 1, 11])                                                                                                                                                                                                                                                                  
Out[5]: tensor([[[1., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0.]]])

In [6]: conv_res = F.conv1d(inputs, kernels, padding=kernel_size//2 + 1, groups=1)  
   ...: conv_res                                                                                                                                                                                                                                                                 
Out[6]: 
tensor([[[0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 0., 0., 0., 0.,
          0.]],

        [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1.,
          1.]]])

In [7]: pad_res = F.pad(conv_res, (1,1), mode='constant', value=0)  
   ...: pad_res                                                                                                                                                                                                                                                                  
Out[7]: 
tensor([[[0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 0., 0., 0.,
          0., 0., 0.]],

        [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1., 1.,
          1., 1., 0.]]])

In [8]: assert pad_res.shape == inputs.shape  

So far, we see that the input that contained original a single “1” has now kernel_size “1”. This is expected.
What if we would like a different kernel per sample? How could we define two different kernels into a same F.conv1d?
I tried the following but sadly, it does not work.
Does anyone know what am I doing wrong?
In [1]: import torch       
   ...: import torch.nn as nn 
   ...: import torch.nn.functional as F                                                                                                                                                                                                                                          

In [2]: inputs1 = [0 for _ in range(20)]     
   ...: inputs1[8] = 1 
   ...: inputs2 = [0 for _ in range(20)] 
   ...: inputs2[13] = 1 
   ...: inputs = torch.FloatTensor([inputs1, inputs2]) 
   ...:                                                                                                                                                                                                                                                                          

In [3]: kernel_size1 = 6 # Let's assume  
   ...: kernel_size2 = 4 
   ...: kernel1 = torch.cat([torch.ones(kernel_size1), torch.zeros(kernel_size1 - 1)]) 
   ...: kernel2 = torch.cat([torch.ones(kernel_size2), torch.zeros(kernel_size2 - 1)]) 
   ...: kernels = torch.cat([kernel1, kernel2], axis=0) 
   ...:                                                                                                                                                                                                                                                                          

In [4]: inputs = inputs.view(-1).unsqueeze(0).unsqueeze(0) 
   ...: inputs                                                                                                                                                                                                                                                                   
Out[4]: 
tensor([[[0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.,
          0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.,
          0., 0., 0., 0., 0., 0.]]])

In [5]: kernels = kernels.unsqueeze(0).unsqueeze(0) 
   ...: kernels                                                                                                                                                                                                                                                                  
Out[5]: 
tensor([[[1., 1., 1., 1., 1., 1., 0., 0., 0., 0., 0., 1., 1., 1., 1., 0., 0.,
          0.]]])

In [6]: conv_res = F.conv1d(inputs, kernels, padding=max(kernel_size1, kernel_size2)//2+1, groups=1)  
   ...: conv_res # Aïe                                                                                                                                                                                                                                                           
Out[6]: 
tensor([[[1., 1., 0., 0., 0., 0., 0., 1., 1., 1., 1., 1., 1., 0., 0., 0., 0.,
          0., 0., 0., 0., 0., 0., 1., 1., 1., 1., 0., 0., 0., 0.]]])

In [7]: conv_res.shape                                                                                                                                                                                                                                                           
Out[7]: torch.Size([1, 1, 31])

Thank you for your help!