PyTorch how to implement disconnection (connections and corresponding gradients are masked)?
I try to implement the following graph. As you can see, the neurons are not fully connected, i.e., the weights are masked and so are their corresponding gradients.



import torch
import numpy as np

x = torch.rand((3, 1))
# tensor([[ 0.8525],
#         [ 0.1509],
#         [ 0.9724]])

weights = torch.rand((2, 3), requires_grad=True)
# tensor([[ 0.3240,  0.0792,  0.6858],
#         [ 0.5248,  0.4565,  0.3625]])

mask = torch.Tensor([[0,1,0],[1,0,1]])
# tensor([[ 0.,  1.,  0.],
#         [ 1.,  0.,  1.]])

mask_weights = weights * mask
# tensor([[ 0.0000,  0.0792,  0.0000],
#         [ 0.5248,  0.0000,  0.3625]])

y = torch.mm(mask_weights, x)
# tensor([[ 0.0120],
#         [ 0.7999]])


This question is originally posted at Pytorch Forum. Note the above way 


  mask_weights = weights * mask


is NOT suitable since corresponding gradients are not 0.

Is there an elegant way to do that please?

Thank you in advance.

Actually, the above method is correct. The disconnections essentially block feed-forward and back-propogation on corresponding connections. In other words, weights and gradients are masked. The codes in question reveal the first while this answer reveals the latter.

mask_weights.register_hook(print)

z = torch.Tensor([[1], [1]])
# tensor([[ 1.],
#         [ 1.]])

out = (y-z).mean()
# tensor(-0.6595)

out.backward()
# tensor([[ 0.1920,  0.1757,  0.0046],
#         [ 0.1920,  0.1757,  0.0046]])

weights.grad
# tensor([[ 0.0000,  0.1757,  0.0000],
#         [ 0.1920,  0.0000,  0.0046]])


As you can see, the gradients of weights are masked automatically.
