Introduction

Setting the k largest values of a tensor to zero is a common masking operation in TensorFlow. The main design choice is whether k should be computed over the whole tensor or independently along an axis, because the indexing strategy changes depending on that requirement.

Core Sections

Global top-k masking

If you want the k largest values in the entire tensor, the simplest approach is to flatten the tensor, find the top indices, zero them out, and then reshape back.

1import tensorflow as tf
2
3x = tf.constant([[1.0, 7.0, 3.0],
4                 [9.0, 2.0, 6.0]])
5
6k = 2
7flat = tf.reshape(x, [-1])
8_, top_indices = tf.math.top_k(flat, k=k)
9
10updated_flat = tf.tensor_scatter_nd_update(
11    flat,
12    indices=tf.expand_dims(top_indices, axis=1),
13    updates=tf.zeros([k], dtype=flat.dtype),
14)
15
16result = tf.reshape(updated_flat, tf.shape(x))
17print(result.numpy())

Output:

[[1. 7. 3.]
 [0. 2. 0.]]

The values 9 and 6 were the two largest across the full tensor, so those positions became zero.

Why top_k plus scatter works well

tf.math.top_k gives you the largest values and their indices efficiently. tf.tensor_scatter_nd_update then lets you replace specific elements without writing Python loops. This keeps the whole operation inside TensorFlow, which matters for performance and graph compatibility.

The general pattern is:

1. choose the values to modify
2. get their indices
3. scatter replacement values at those indices

That pattern is useful beyond zeroing too. You can replace the top elements with a threshold, a constant, or a learned value if your workflow needs it.

Zeroing top values row by row

Sometimes "set the k largest elements to zero" really means "do this independently for each row." In that case, a global flatten is wrong because it mixes all rows together. For per-row behavior, use top_k along the last dimension and build row-aware scatter indices.

1import tensorflow as tf
2
3x = tf.constant([[1.0, 7.0, 3.0],
4                 [9.0, 2.0, 6.0]])
5
6k = 1
7_, col_indices = tf.math.top_k(x, k=k)
8
9row_indices = tf.repeat(tf.range(tf.shape(x)[0]), repeats=k)
10col_indices = tf.reshape(col_indices, [-1])
11
12scatter_indices = tf.stack([row_indices, col_indices], axis=1)
13updates = tf.zeros([tf.shape(scatter_indices)[0]], dtype=x.dtype)
14
15result = tf.tensor_scatter_nd_update(x, scatter_indices, updates)
16print(result.numpy())

Output:

[[1. 0. 3.]
 [0. 2. 6.]]

Now each row lost only its own largest element.

A masking alternative

For some workflows, it is easier to build a boolean or numeric mask than to scatter updates manually. This is especially useful if you later want to reuse the mask for other operations.

1import tensorflow as tf
2
3x = tf.constant([1.0, 7.0, 3.0, 9.0, 2.0, 6.0])
4k = 2
5
6values, indices = tf.math.top_k(x, k=k)
7mask = tf.ones_like(x)
8mask = tf.tensor_scatter_nd_update(
9    mask,
10    tf.expand_dims(indices, 1),
11    tf.zeros([k], dtype=x.dtype),
12)
13
14result = x * mask
15print(result.numpy())

That approach is especially handy when you want to preserve differentiability through the surviving elements while making the zeroing rule explicit.

What about ties

If several elements share the same value around the cutoff, top_k still returns exactly k positions. That means some equal-valued elements may be zeroed while others remain, depending on their returned index order. If tie handling matters semantically, define that behavior up front rather than assuming all equal maxima will be treated identically.

Common Pitfalls

• Flattening the tensor when the real requirement is to zero the top values per row or per batch item.
• Using Python loops over tensor elements instead of TensorFlow ops such as top_k and scatter updates.
• Forgetting that scatter indices must have the right rank and shape for the target tensor.
• Assuming tied values will all be zeroed even though top_k returns exactly k indices.
• Rebuilding large masks unnecessarily when a direct scatter update would be simpler and cheaper.

Summary

• Use tf.math.top_k to identify the largest values and their indices.
• For global top-k, flatten the tensor, update the selected positions, and reshape back.
• For per-row or per-axis behavior, build scatter indices that preserve that structure.
• 'tf.tensor_scatter_nd_update is the usual TensorFlow tool for zeroing selected elements.'
• Be explicit about whether k applies globally or within each slice of the tensor.