How can I set the diagonal of an N-dim tensor to 0 along given dims?

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(3,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]]])

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

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.

Is there a way to mask a given diagonal over 2 arbitrary dimensions for a 3D-tensor?

Solution

You can do this with a for loop over the sub-tensors:

# across dim0
for i in range(data.size(0)):
    data[i].fill_diagonal_(0)

If you need to perform this over an arbitrary two dimensions of a 3d tensor, simply apply the fill to the appropriate slices:

# across dim1
for i in range(data.size(1)):
    data[:,i].fill_diagonal_(0)

# across dim2
for i in range(data.size(2)):
    data[:,:,i].fill_diagonal_(0)