Introduction

Zeroing the diagonal of a tensor is a common cleanup step when diagonal entries represent self-links, self-similarity, or values you want to exclude from later computation. The exact code depends on whether you are using a mutable array library such as NumPy or an immutable tensor library such as TensorFlow.

Zero the Diagonal in TensorFlow

In TensorFlow, tensors are immutable, so you do not modify them in place. Instead, create a new tensor with a replaced diagonal by using tf.linalg.set_diag.

1import tensorflow as tf
2
3x = tf.constant(
4    [
5        [1.0, 2.0, 3.0],
6        [4.0, 5.0, 6.0],
7        [7.0, 8.0, 9.0],
8    ]
9)
10
11zeroed = tf.linalg.set_diag(x, tf.zeros([3], dtype=x.dtype))
12
13print(zeroed.numpy())

The diagonal argument must match the diagonal length of the last two dimensions. For a square 3 x 3 matrix, that means a vector of length 3.

Batched Tensors

If your tensor contains a batch of matrices, tf.linalg.set_diag still works. You just need one diagonal vector per matrix in the batch.

1import tensorflow as tf
2
3batch = tf.constant(
4    [
5        [[1.0, 2.0], [3.0, 4.0]],
6        [[5.0, 6.0], [7.0, 8.0]],
7    ]
8)
9
10diagonals = tf.zeros([2, 2], dtype=batch.dtype)
11zeroed_batch = tf.linalg.set_diag(batch, diagonals)
12
13print(zeroed_batch.numpy())

Here the input shape is (2, 2, 2), which means two 2 x 2 matrices. The diagonal tensor shape is (2, 2), one row of diagonal values for each matrix.

Non-Square Matrices

The diagonal length is the smaller of the row count and column count. For example, a 3 x 5 matrix has a diagonal length of 3. You still use tf.linalg.set_diag, but the replacement vector must follow that shorter length.

1rect = tf.constant(
2    [
3        [1.0, 2.0, 3.0, 4.0, 5.0],
4        [6.0, 7.0, 8.0, 9.0, 10.0],
5        [11.0, 12.0, 13.0, 14.0, 15.0],
6    ]
7)
8
9zeroed_rect = tf.linalg.set_diag(rect, tf.zeros([3], dtype=rect.dtype))
10print(zeroed_rect.numpy())

That detail matters because shape errors around diagonal length are one of the most common reasons this operation fails.

NumPy Alternative

If you are working with NumPy arrays instead of TensorFlow tensors, the operation can be done in place:

1import numpy as np
2
3arr = np.array(
4    [
5        [1.0, 2.0, 3.0],
6        [4.0, 5.0, 6.0],
7        [7.0, 8.0, 9.0],
8    ]
9)
10
11np.fill_diagonal(arr, 0.0)
12print(arr)

This difference is important. In TensorFlow, you create a new tensor. In NumPy, fill_diagonal mutates the existing array.

Why This Operation Shows Up Often

Zeroing diagonals is common in several domains:

• Graph adjacency matrices when self-loops should be removed
• Similarity matrices when self-similarity should not dominate top matches
• Pairwise distance computations when diagonal entries are trivial

Because the operation is so common, it is worth using a library function instead of hand-writing index loops unless you have a very specific performance reason.

Common Pitfalls

The biggest pitfall in TensorFlow is expecting the original tensor to change in place. It will not. tf.linalg.set_diag returns a new tensor and leaves the original unchanged.

Another common mistake is passing the wrong diagonal shape, especially for batched tensors. The diagonal replacement must align with the last two dimensions of the input tensor.

It is also easy to forget dtype matching. If the input tensor is floating point and the diagonal vector is integer by default, TensorFlow may complain or force an unwanted cast. Creating the zeros with dtype=x.dtype avoids that mismatch.

Finally, if you try to solve this with manual index assignment copied from NumPy examples, remember that TensorFlow tensors are not normal mutable Python arrays. Use the tensor-specific API instead.

Summary

• In TensorFlow, use tf.linalg.set_diag and store the returned tensor.
• For batched matrices, provide one diagonal vector per matrix.
• For non-square matrices, the diagonal length is the smaller of the last two dimensions.
• In NumPy, np.fill_diagonal is a simple in-place alternative.
• Watch for shape and dtype mismatches when building the diagonal replacement values.