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

1

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! :slight_smile:

2 Likes
2

Hi Alpha!

In the case that you want to zero out the 1-2 diagonal of all three
dim = 0 slices, you can use broadcasting, together with element-wise
multiplication by a two-dimensional zero-diagonal matrix:

>>> data = torch.randn ((3, 3, 3))
>>> print (data)
tensor([[[-0.0236, -0.3462, -0.3220],
         [ 0.3032,  1.3593, -1.1444],
         [-0.0366, -0.4146,  0.7503]],

        [[ 1.3535,  1.0488,  1.1539],
         [-0.2239,  0.3035, -1.6576],
         [-0.3756,  1.4249,  0.6928]],

        [[-0.0917,  2.4133,  0.0379],
         [-0.0485, -0.5681, -0.3200],
         [-0.5743, -2.4652, -0.2874]]])
>>> dataz = (1 - torch.eye (3)) * data
>>> print (dataz)
tensor([[[-0.0000, -0.3462, -0.3220],
         [ 0.3032,  0.0000, -1.1444],
         [-0.0366, -0.4146,  0.0000]],

        [[ 0.0000,  1.0488,  1.1539],
         [-0.2239,  0.0000, -1.6576],
         [-0.3756,  1.4249,  0.0000]],

        [[-0.0000,  2.4133,  0.0379],
         [-0.0485, -0.0000, -0.3200],
         [-0.5743, -2.4652, -0.0000]]])

You can either transpose() your data matrix so that your desired
dimensions become your 1-2 dimensions, use broadcasting, as above,
and then transpose() it back, or you can build a three-dimensional
“mask” tensor using a two-dimensional zero-diagonal matrix with
repeat(), transpose it to put the zero-diagonals along the desired
dimensions, and then element-wise multiply your three-dimensional
“mask” tensor with your three-dimensional data tensor:

>>> mask01 = (1 - torch.eye (3)).repeat (3, 1, 1).transpose (0, 2)
>>> print (mask01)
tensor([[[0., 0., 0.],
         [1., 1., 1.],
         [1., 1., 1.]],

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

        [[1., 1., 1.],
         [1., 1., 1.],
         [0., 0., 0.]]])
>>> dataz01 = mask01 * data
>>> print (dataz01)
tensor([[[-0.0000, -0.0000, -0.0000],
         [ 0.3032,  1.3593, -1.1444],
         [-0.0366, -0.4146,  0.7503]],

        [[ 1.3535,  1.0488,  1.1539],
         [-0.0000,  0.0000, -0.0000],
         [-0.3756,  1.4249,  0.6928]],

        [[-0.0917,  2.4133,  0.0379],
         [-0.0485, -0.5681, -0.3200],
         [-0.0000, -0.0000, -0.0000]]])

You could apply either of the above techniques multiple times, once for
each of your pairs of dimensions, or, if you reuse the same set of pairs
of dimensions, you could make a “multi-pair” mask tensor by multiplying
together the respective “single-pair” mask tensors, and use that
multi-pair mask tensor every time you need to zero out that particular
set of diagonals.

Best.

K. Frank

1 Like
3

Hi KFrank,

Thanks for the clear explanation!

Cheers,

4

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