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

AlphaBetaGamma96: <p>Hi All,</p>
<p>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 <code>[N,N,N]</code> and I wanted to set the diagonal along dim=1,2 equal to 0? How exactly could that be done?</p>
<p>I tried using <code>fill_diagonal_</code> but that only does the k-th diagonal element for each sub-array, i.e,</p>
<pre><code class="lang-auto">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.]]])

</code></pre>
<p>whereas I would want the entire diagonal for each sub-matrix to be equal to 0 here. So, the desired outcome would be,</p>
<pre><code class="lang-auto">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.]]])
</code></pre>
<p>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.</p>
<p>In short, is there a way to mask a given diagonal over 2 arbitrary dimensions for a 3D-tensor?</p>
<p>Thank you in advance! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

KFrank: <p>Hi Alpha!</p>
<aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="1" data-topic="107448" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p>I’m trying to figure out a way to set the diagonal of a 3-dimensional Tensor (along 2 given dims) equal to 0.</p>
</blockquote>
</aside>
<p>In the case that you want to zero out the 1-2 diagonal of all three<br>
dim = 0 slices, you can use broadcasting, together with element-wise<br>
multiplication by a two-dimensional zero-diagonal matrix:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; data = torch.randn ((3, 3, 3))
&gt;&gt;&gt; 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]]])
&gt;&gt;&gt; dataz = (1 - torch.eye (3)) * data
&gt;&gt;&gt; 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]]])
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>In short, is there a way to mask a given diagonal over 2 arbitrary dimensions for a 3D-tensor?</p>
</blockquote>
</aside>
<p>You can either <code>transpose()</code> your data matrix so that your desired<br>
dimensions become your 1-2 dimensions, use broadcasting, as above,<br>
and then <code>transpose()</code> it back, or you can build a three-dimensional<br>
“mask” tensor using a two-dimensional zero-diagonal matrix with<br>
<code>repeat()</code>, transpose it to put the zero-diagonals along the desired<br>
dimensions, and then element-wise multiply your three-dimensional<br>
“mask” tensor with your three-dimensional <code>data</code> tensor:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; mask01 = (1 - torch.eye (3)).repeat (3, 1, 1).transpose (0, 2)
&gt;&gt;&gt; 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.]]])
&gt;&gt;&gt; dataz01 = mask01 * data
&gt;&gt;&gt; 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]]])
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>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.</p>
</blockquote>
</aside>
<p>You could apply either of the above techniques multiple times, once for<br>
each of your pairs of dimensions, or, if you reuse the same set of pairs<br>
of dimensions, you could make a “multi-pair” mask tensor by multiplying<br>
together the respective “single-pair” mask tensors, and use that<br>
multi-pair mask tensor every time you need to zero out that particular<br>
set of diagonals.</p>
<p>Best.</p>
<p>K. Frank</p>

AlphaBetaGamma96: <p>Hi KFrank,</p>
<p>Thanks for the clear explanation!</p>
<p>Cheers,</p>

iacob: <p>You can do this with a <a href="https://stackoverflow.com/questions/66765173/how-can-i-set-the-diagonal-of-an-n-dim-tensor-to-0-along-given-dims/66765242#66765242" rel="noopener nofollow ugc">for loop</a> over the sub-tensors:</p>
<pre><code class="lang-python"># across dim0
for i in range(data.size(0)):
    data[i].fill_diagonal_(0)
</code></pre>
<p>If you need to perform this over an arbitrary two dimensions of a 3d tensor, simply apply the fill to the appropriate slices:</p>
<pre><code class="lang-python"># across dim1
for i in range(data.size(1)):
    data[:,i].fill_diagonal_(0)
</code></pre>
<pre><code class="lang-python"># across dim2
for i in range(data.size(2)):
    data[:,:,i].fill_diagonal_(0)
</code></pre>