I have a tensor T of shape (b, r)
I want to do an operation for each (r), in a way that it gets parallelized by the GPU
The naive implementation, in numpy for simplicity, would look something like:
T_dash = np.array([(T[i] - np.max(T[i]) for i in range(T.size[0])])
What would be the best way to do this?
There's a new vmap function available (in the master branch at the time of writing, experimental) that will help do batch operations, where you define the operation to be performed for each element.
vmap can be helpful in hiding batch dimensions. In your case, it goes something like
def each_elem_fn(tensor):
return tensor - np.max(tensor)
torch.vmap(each_elem_fn)(batch_of_elems)
Note: As @jodag mentioned in the comment, for simply subtracting max you can use
T = T - T.max(dim=1, keepdim=True)[0]
however, if you want to generalize it to any custom function you shall use vmap.