Hi All,
I’m trying to figure out a way to set the diagonal of a 3-dimensional Tensor (along 2 given dims) equal to 0. An example of this would be, let’s say I have a Tensor of shape [N,N,N] and I wanted to set the diagonal along dim=1,2 equal to 0? How exactly could that be done?
I tried using fill_diagonal_ but that only does the k-th diagonal element for each sub-array, i.e,
data = torch.ones(4,4,4)
 
data.fill_diagonal_(0)
tensor([[[0., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.]],

        [[1., 1., 1., 1.],
         [1., 0., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.]],

        [[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 0., 1.],
         [1., 1., 1., 1.]],

        [[1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 1.],
         [1., 1., 1., 0.]]])


whereas I would want the entire diagonal for each sub-matrix to be equal to 0 here. So, the desired outcome would be,
tensor([[[0., 1., 1., 1.],
         [1., 0., 1., 1.],
         [1., 1., 0., 1.],
         [1., 1., 1., 0.]],

        [[0., 1., 1., 1.],
         [1., 0., 1., 1.],
         [1., 1., 0., 1.],
         [1., 1., 1., 0.]],

        [[0., 1., 1., 1.],
         [1., 0., 1., 1.],
         [1., 1., 0., 1.],
         [1., 1., 1., 0.]],

        [[0., 1., 1., 1.],
         [1., 0., 1., 1.],
         [1., 1., 0., 1.],
         [1., 1., 1., 0.]]])

Secondly, the reason I state for a given pair of dimension is, I need to repeat this `zeroing’ along 2 different pairs of dimensions (e.g. dim=(1,2) then dim=(0,1) ) to get the required masking I need.
In short, is there a way to mask a given diagonal over 2 arbitrary dimensions for a 3D-tensor?
Thank you in advance!