Introduction

Tensor slicing is a core operation in TensorFlow pipelines, model debugging, and feature engineering. You can slice tensors with Python-style syntax for simple ranges, then switch to specialized ops for dynamic or sparse index patterns. This guide covers both approaches with runnable examples.

Basic Slice Syntax

For contiguous ranges, TensorFlow supports familiar indexing syntax.

1import tensorflow as tf
2
3x = tf.constant([
4    [10, 11, 12, 13],
5    [20, 21, 22, 23],
6    [30, 31, 32, 33],
7], dtype=tf.int32)
8
9print(x[0])         # first row
10print(x[:, 1:3])    # all rows, columns 1 and 2
11print(x[::2, :])    # every second row

Rules follow Python indexing semantics, including start, stop, and step behavior.

Use tf.gather for Index Lists

When you need non-contiguous indices, tf.gather is clearer than manual slicing.

1row_ids = tf.constant([2, 0])
2selected_rows = tf.gather(x, row_ids, axis=0)
3print(selected_rows)
4
5col_ids = tf.constant([3, 1])
6selected_cols = tf.gather(x, col_ids, axis=1)
7print(selected_cols)

tf.gather is ideal when index positions come from model logic or preprocessing steps.

Use tf.gather_nd for Coordinate-Based Extraction

If you need values from specific coordinates across multiple dimensions, use tf.gather_nd.

1coords = tf.constant([
2    [0, 1],
3    [1, 3],
4    [2, 2],
5])
6
7values = tf.gather_nd(x, coords)
8print(values)  # [11, 23, 32]

Each coordinate points to one element. This is useful for sparse lookup patterns.

Dynamic Slicing with tf.slice

tf.slice works well when start offsets and lengths are computed at runtime.

1start = tf.constant([1, 1])
2size = tf.constant([2, 2])
3block = tf.slice(x, begin=start, size=size)
4print(block)

Unlike Python syntax, tf.slice accepts tensors for begin and size, which is useful inside graph execution.

Boolean Filtering with tf.boolean_mask

For predicate-based selection, create a mask and filter rows or elements.

1scores = tf.constant([0.2, 0.8, 0.6, 0.1], dtype=tf.float32)
2labels = tf.constant([0, 1, 1, 0], dtype=tf.int32)
3mask = scores > 0.5
4
5high_scores = tf.boolean_mask(scores, mask)
6high_labels = tf.boolean_mask(labels, mask)
7
8print(high_scores)
9print(high_labels)

This pattern is common in post-processing predictions.

Integrating Slicing into tf.data

Use slicing ops in input pipelines so logic remains vectorized and reproducible.

1def transform(features, target):
2    # Keep only selected feature indices
3    features = tf.gather(features, [0, 2, 4], axis=-1)
4    return features, target
5
6features = tf.random.normal((100, 6))
7targets = tf.random.uniform((100,), maxval=2, dtype=tf.int32)
8
9ds = tf.data.Dataset.from_tensor_slices((features, targets))
10ds = ds.map(transform, num_parallel_calls=tf.data.AUTOTUNE).batch(16)
11
12for batch_x, batch_y in ds.take(1):
13    print(batch_x.shape, batch_y.shape)

Embedding indexing in map avoids ad-hoc preprocessing outside the training graph.

Shape Safety and Debugging Techniques

Indexing bugs are often shape bugs. Add explicit checks so failures happen early.

1def safe_take_columns(batch):
2    tf.debugging.assert_rank(batch, 2)
3    tf.debugging.assert_greater_equal(tf.shape(batch)[1], 5)
4    return tf.gather(batch, [0, 2, 4], axis=1)
5
6sample = tf.random.normal((8, 6))
7out = safe_take_columns(sample)
8print(out.shape)

For model debugging, print shapes after each slice step during development. In production input pipelines, prefer assertions and unit tests over frequent logging.

Choosing the Right Operation

Use simple slice syntax when range boundaries are static and easy to read. Move to tf.gather when selection is index-driven, and to tf.gather_nd when you need coordinate lookup across dimensions. This keeps code clear and helps reviewers reason about expected output shape quickly.

Common Pitfalls

• Mixing Python lists and tensor indices in graph-heavy code paths, leading to conversion overhead or shape surprises.
• Using advanced indexing assumptions from NumPy that do not map exactly to TensorFlow ops.
• Forgetting bounds checks when index tensors come from model outputs.
• Applying tf.boolean_mask on the wrong axis and getting unexpected flattened output.
• Repeating costly slicing in Python loops instead of vectorized dataset mapping.

Summary

• Use Python-style slicing for simple contiguous ranges.
• Use tf.gather for index lists and tf.gather_nd for coordinate-based lookup.
• Use tf.slice when offsets and lengths are dynamic tensors.
• Use tf.boolean_mask for predicate-driven filtering.
• Keep slicing inside tf.data pipelines for consistency and performance.