Search⌘ K
AI Features

Tensor Indexing

Explore tensor indexing techniques in TensorFlow to retrieve final word predictions from sequences using tf.gather and tf.gather_nd. Understand how to create indices for batch sequences and apply these methods within LSTM-based language models.

Chapter Goals:

  • Extract the word predictions for the final time step of each sequence

A. Gather functions

When generating word predictions for a batch of sequences, we only want to retrieve the word IDs for each sequence's final time step. If this were regular Python (or NumPy), we could do this through simple list indexing or slicing. However, in TensorFlow we need to use gather functions to retrieve data at specific locations in the tensor.

There are two main gather functions: tf.gather and tf.gather_nd. Both take in the same two required arguments:

  • ...