# Efficiently computing a batch of results given a batch assignment vector and series of corresponding matrices

I have a 1D tensor of tokens that belong to different batches. The batch sizes here are uneven. Each batch needs to be multiplied with a corresponding weight matrix. My current approach is using a batch pointer vector and a series of distinct weight matrices corresponding to the unique pointers along with a for loop. I want to efficiently compute a result of shape [num_tokens, output_dim], where each weight matrix has shape [input_dim, output_dim]. I also pad the inputs to a multiple of 8 for harnessing NVIDIA Tensor Cores. Here's an example:

import torch

# shape [num_tokens,]
input_dim, output_dim = 4, 8
ptr = torch.tensor([0, 1, 2, 2, 3, 3, -1, -1]) # -1 means padding
features = torch.randn(ptr.shape, input_dim)
weights = [torch.randn(input_dim, output_dim) for _ in range(4)]

unique = torch.unique(ptr, sorted=False, return_inverse=False, return_counts=False)
unique = unique[unique != -1] # ignore padding

results = []

for i in unique:
split = features[ptr == i, :]
# pad each split to multiple of 8 for NVIDIA A100
# repeat pad embedding to desired size
torch.empty((-split.size(0)) % 8, split.size(-1))
.uniform_()
.to(split.device)
)
torch.bool
)

# forward pass
# strip padding so I can create a 2D result tensor of correct dimension again
# we can now pad results to match the shape of ptr