Python: Vectorización de Cálculo Implementado utilizando el Enfoque Iterativo

Ya que se trata de tres dimensiones n, my c (en minúsculas para evitar la ambigüedad), puede ser útil para cambiar la forma de todos sus tensores para (n, m, c)mediante la replicación de sus valores sobre la dimensión que falta (por ejemplo, M(m, c) se convierte en M(n, m, c)).

Sin embargo, usted puede saltar el explícito de replicación y el uso de la radiodifusión, por lo que es suficiente para que estire la dimensión que falta (por ejemplo, M(m, c) se convierte en M(1, m, c).

Teniendo en cuenta estas consideraciones, la vectorización de su código se hace de la siguiente manera

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)

He aquí una pequeña prueba para comprobar la validez de la solución

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))

Siéntase libre de probar en otros casos así!