Introduction

In TensorFlow, the standard pattern is to build a boolean mask, use tf.where to get the coordinates of nonzero elements, and then use tf.gather_nd or tf.boolean_mask to get the values. Which combination you choose depends on whether you need both indices and values or only the filtered values.

Use tf.where to get nonzero indices

tf.where returns the coordinates where a condition is true. For nonzero detection, the condition is usually tensor != 0.

1import tensorflow as tf
2
3x = tf.constant([
4    [0, 5, 0],
5    [7, 0, 9],
6], dtype=tf.int32)
7
8indices = tf.where(x != 0)
9print(indices.numpy())

Output:

[[0 1]
 [1 0]
 [1 2]]

Each row in the result is an index tuple. In this two-dimensional example, [0, 1] means row 0, column 1.

Use tf.gather_nd to retrieve the matching values

If you already have the coordinates and also want the values, tf.gather_nd is the clean companion function.

1import tensorflow as tf
2
3x = tf.constant([
4    [0, 5, 0],
5    [7, 0, 9],
6], dtype=tf.int32)
7
8indices = tf.where(x != 0)
9values = tf.gather_nd(x, indices)
10
11print(indices.numpy())
12print(values.numpy())

Output:

1[[0 1]
2 [1 0]
3 [1 2]]
4[5 7 9]

This is the most direct answer when you need both pieces of information in matching order.

Use tf.boolean_mask when you only care about values

If you do not need coordinates, tf.boolean_mask is shorter and easier to read.

1import tensorflow as tf
2
3x = tf.constant([
4    [0, 5, 0],
5    [7, 0, 9],
6], dtype=tf.int32)
7
8values_only = tf.boolean_mask(x, x != 0)
9print(values_only.numpy())

That returns a one-dimensional tensor of the selected values. It is ideal for filtering, but it does not preserve the full coordinate structure by itself.

Understand the difference for higher-rank tensors

The same pattern works for three-dimensional and higher-dimensional tensors. The only difference is that each index row returned by tf.where becomes longer.

1import tensorflow as tf
2
3x = tf.constant([
4    [[0, 1], [2, 0]],
5    [[0, 0], [3, 4]],
6], dtype=tf.int32)
7
8indices = tf.where(x != 0)
9values = tf.gather_nd(x, indices)
10
11print(indices.numpy())
12print(values.numpy())

Here, each index has three numbers because the tensor rank is 3.

Be clear about zeros versus near-zero floats

For integer tensors, x != 0 is usually exactly what you want. For floating-point tensors, exact comparison to zero can be too strict if your data comes from computation rather than literal constants.

In those cases, compare against a tolerance.

1import tensorflow as tf
2
3x = tf.constant([0.0, 1e-8, 0.2, -0.3], dtype=tf.float32)
4mask = tf.abs(x) > 1e-6
5indices = tf.where(mask)
6values = tf.boolean_mask(x, mask)
7
8print(indices.numpy())
9print(values.numpy())

That avoids treating tiny numerical noise as meaningful nonzero data.

Use sparse representations only when the whole workflow is sparse

If you repeatedly do this operation because the tensor is mostly zeros, it may be a sign that a sparse representation is more appropriate. TensorFlow supports sparse tensors, but converting just for one small operation is usually not worth the added complexity.

So use tf.where and masking for ordinary tensor filtering, and consider sparse structures only when the surrounding pipeline is already sparse-aware.

Common Pitfalls

• Using tf.boolean_mask when you actually need coordinates as well as values.
• Comparing floating-point tensors to exact zero when a tolerance would be more appropriate.
• Forgetting that tf.where(condition) returns indices, not the selected values themselves.
• Assuming the result of tf.boolean_mask keeps the original tensor shape.
• Reaching for sparse tensors too early when simple masking is enough.

Summary

• Use tf.where(x != 0) to get the indices of nonzero elements.
• Use tf.gather_nd with those indices when you also need the values.
• Use tf.boolean_mask when you only need the filtered nonzero values.
• For floating-point tensors, consider a tolerance instead of exact zero comparison.
• The same pattern works for tensors of any rank; the index tuples just get longer.