Tensorflow Get indices of array rows which are zero
Master System Design with Codemia
Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.
Introduction
To find rows that are entirely zero in TensorFlow, the standard pattern is to compare values to zero, reduce across each row, and then extract the matching indices. The steps are short, but axis choice and output shape often trip people up the first time.
Use tf.equal, tf.reduce_all, and tf.where
For a 2D tensor, the core idea is:
- compare every element with zero
- reduce across columns to decide whether each row is all zero
- use
tf.whereto get the row positions
The boolean mask is one value per row. tf.where then returns the positions where that mask is True.
Why axis=1 Is Correct for Rows
For a 2D tensor, axis=1 means “reduce across columns within each row”. That is exactly what you want when asking whether a row is all zero.
If you accidentally use axis=0, you are checking whether each column is entirely zero, which answers a different question.
Flatten the Result If You Need a 1D Index Tensor
tf.where returns a 2D index tensor, so the row indices come out looking like [[0], [2]]. If you want a flat result such as [0, 2], slice the first column:
This shape detail matters because later code may expect a simple 1D tensor for indexing or logging.
Get the Zero Rows Themselves
If the next step is to inspect or drop those rows, tf.boolean_mask is often more convenient than just keeping their positions:
This is useful for padded batches, sparse data cleanup, or debugging model inputs that unexpectedly contain empty rows.
Floating-Point Tensors Need Tolerance
Exact equality to zero is often too strict for floating-point tensors because small rounding noise can remain after computation. In those cases, compare against a threshold:
This is a better pattern for outputs from neural network layers or numerical preprocessing pipelines, where “effectively zero” often matters more than exact bitwise zero.
Wrap It in a Helper Function
If the operation appears repeatedly, turn it into a utility:
Small helpers reduce duplicate mask-building code and make later tensor pipelines easier to read.
Common Pitfalls
- Using
axis=0when the goal is to test rows rather than columns. - Forgetting that
tf.wherereturns a 2D tensor of indices, not a flat vector. - Confusing rows that contain at least one zero with rows whose every value is zero.
- Using exact equality on floating-point data when a tolerance test is safer.
- Applying the same axis logic to higher-rank tensors without first defining what a “row” means in that shape.
Summary
- The standard TensorFlow pattern is
tf.equal, thentf.reduce_all, thentf.where. - For 2D tensors,
axis=1checks whether each row is entirely zero. - Slice
[:, 0]fromtf.wherewhen you want a flat vector of row indices. - Use
tf.boolean_maskif you want the rows themselves instead of just their positions. - Prefer a tolerance-based check for floating-point tensors that may contain tiny numerical noise.

