Python: Векторизация расчета, реализованная с использованием итеративного подхода

Поскольку вы имеете дело с тремя измерениями n, m, и c (в нижнем регистре, чтобы избежать двусмысленности), может быть полезно изменить форму всех ваших тензоров на (n, m, c), путем репликации их значений по отсутствующему измерению (например M(m, c) становится M(n, m, c)).

Однако вы можете пропустить явную репликацию и использовать широковещательную передачу, поэтому достаточно отменить запрос на недостающее измерение (например M(m, c) становится M(1, m, c).

Учитывая эти соображения, векторизация вашего кода становится следующей

cond = y.unsqueeze(2) == torch.arange(M.size(dim=1)).unsqueeze(0)
pos = ds.unsqueeze(2) * M.unsqueeze(1) * cond
neg = ds.unsqueeze(2) * M.unsqueeze(1).sum(dim=0, keepdim=True) * ~cond
grad += (pos - neg).sum(dim=0)

Вот небольшой тест для проверки правильности решения

import torch

n, m, c = 11, 5, 7

y = torch.randint(c, size=(n, 1))
ds = torch.rand(n, m)
M = torch.randint(2, size=(n, c))
grad = torch.rand(m, c)


def slow_grad(y, ds, M, grad):
    for i in range(M.size(dim=1)):
        for j in range(ds.size(dim=0)):
            if y[j] == i:
                grad[:,i] = grad[:,i] - (ds[j,:].T * sum(M[j,:]))
            else:
                if M[j,i] > 0:
                    grad[:,i] = grad[:,i] + ds[j,:].T
    return grad


def fast_grad(y, ds, M, grad):
    cond = y.unsqueeze(2) == torch.arange(M.size(dim=1)).unsqueeze(0)
    pos = ds.unsqueeze(2) * M.unsqueeze(1) * cond
    neg = ds.unsqueeze(2) * M.unsqueeze(1).sum(dim=0, keepdim=True) * ~cond
    grad += (pos - neg).sum(dim=0)
    return grad
  
# Assert equality of all elements function outputs, throws an exception if false
assert torch.all(slow_grad(y, ds, M, grad) == fast_grad(y, ds, M, grad))

Не стесняйтесь тестировать и на других случаях!