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

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?

2 Likes

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

Hi KFrank,

Thanks for the clear explanation!

Cheers,

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