TensorFlow
zero rows
array indices
machine learning
data manipulation

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.where to get the row positions
python
1import tensorflow as tf
2
3x = tf.constant([
4    [0, 0, 0],
5    [1, 0, 0],
6    [0, 0, 0],
7    [2, 3, 4],
8], dtype=tf.int32)
9
10row_is_zero = tf.reduce_all(tf.equal(x, 0), axis=1)
11indices = tf.where(row_is_zero)
12
13print(row_is_zero.numpy())
14print(indices.numpy())

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.

python
1import tensorflow as tf
2
3x = tf.constant([
4    [0, 0, 0],
5    [1, 0, 0],
6    [0, 0, 0],
7])
8
9by_columns = tf.reduce_all(tf.equal(x, 0), axis=0)
10by_rows = tf.reduce_all(tf.equal(x, 0), axis=1)
11
12print("axis=0:", by_columns.numpy())
13print("axis=1:", by_rows.numpy())

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:

python
1import tensorflow as tf
2
3x = tf.constant([
4    [0, 0],
5    [5, 6],
6    [0, 0],
7])
8
9mask = tf.reduce_all(tf.equal(x, 0), axis=1)
10flat_indices = tf.where(mask)[:, 0]
11
12print(flat_indices.numpy())

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:

python
1import tensorflow as tf
2
3x = tf.constant([
4    [0, 0, 0],
5    [1, 2, 3],
6    [0, 0, 0],
7])
8
9row_is_zero = tf.reduce_all(tf.equal(x, 0), axis=1)
10zero_rows = tf.boolean_mask(x, row_is_zero)
11
12print(zero_rows.numpy())

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:

python
1import tensorflow as tf
2
3x = tf.constant([
4    [0.0, 0.0, 0.0],
5    [1e-8, 0.0, 0.0],
6    [2.0, 3.0, 4.0],
7], dtype=tf.float32)
8
9eps = 1e-7
10near_zero_rows = tf.reduce_all(tf.abs(x) < eps, axis=1)
11indices = tf.where(near_zero_rows)[:, 0]
12
13print(indices.numpy())

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:

python
1import tensorflow as tf
2
3
4def zero_row_indices(x: tf.Tensor) -> tf.Tensor:
5    mask = tf.reduce_all(tf.equal(x, 0), axis=1)
6    return tf.where(mask)[:, 0]
7
8
9x = tf.constant([
10    [0, 0],
11    [9, 1],
12    [0, 0],
13])
14
15print(zero_row_indices(x).numpy())

Small helpers reduce duplicate mask-building code and make later tensor pipelines easier to read.

Common Pitfalls

  • Using axis=0 when the goal is to test rows rather than columns.
  • Forgetting that tf.where returns 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, then tf.reduce_all, then tf.where.
  • For 2D tensors, axis=1 checks whether each row is entirely zero.
  • Slice [:, 0] from tf.where when you want a flat vector of row indices.
  • Use tf.boolean_mask if 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.

Course illustration
Course illustration

All Rights Reserved.