How to properly implement these 1D convolutions and summations?

How can I properly implement the convolution and summation as shown in the example below?

Lets be given a PyTorch tensor of signals of size (batch_size, num_signals, signal_length), i.e. each batch contains several signals. For each batch, I want to convolve the i-th signal with the i-th kernel, and sum all of these convolutions. The result should be of shape (batch_size, 1, signal_length)

The code snipped below should to the right thing, however I want to avoid unnecessary loops.

Thanks for any help!


import torch
import torch.nn as nn
import torch.nn.functional as F

batch_size = 16
num_signals = 5
signal_length = 32
num_kernels = 5
kernel_size = 32

# Random batch of signals and kernels
signals = torch.randn(batch_size, num_signals, signal_length)
kernels = torch.randn(num_kernels, 1, kernel_size)

out = torch.zeros(signals.shape[0], 1, signal_length)
for b in range(signals.shape[0]):
    for k in range(num_signals):
        out[b, 0, :] += F.conv1d(
            signals[b, k].unsqueeze(0).unsqueeze(0),
            kernels[k].unsqueeze(0),
            padding="same"
        ).squeeze()

print("Output shape:", out.shape)