{"id":1522,"date":"2023-12-23T18:00:10","date_gmt":"2023-12-23T18:00:10","guid":{"rendered":"https:\/\/www.philippeadjiman.com\/blog\/?p=1522"},"modified":"2025-07-18T13:13:44","modified_gmt":"2025-07-18T13:13:44","slug":"deep-learning-gymnastics-tensor-indexing","status":"publish","type":"post","link":"https:\/\/philippeadjiman.com\/blog\/2023\/12\/23\/deep-learning-gymnastics-tensor-indexing\/","title":{"rendered":"Deep Learning Gymnastics #2: Tensor Indexing"},"content":{"rendered":"\n

Welcome to the second episode of the Deep Learning Gymnastics<\/a> series. Hope you’re in good shape. Get warmed up. We start.<\/p>\n\n\n\n

Today, we’ll talk about a simple yet important and powerful aspect of tensor manipulations: tensor indexing. <\/p>\n\n\n\n

Batches and embeddings motivating example<\/h2>\n\n\n\n

At the heart of any modern deep learning model, you’ll most often deal with batches and embeddings. <\/p>\n\n\n\n

Batches?<\/strong> Below is a toy example of what a batch from a training set could look like:<\/p>\n\n\n\n

\"\"<\/a><\/figure>\n\n\n\n

The numbers represent an index in a vocabulary of size N, representing any kind of entity. This could be letters or words in a language model, a movie in a recommender system, a segment on a map in an ETA model, or ads in an ad Network. <\/p>\n\n\n\n

For the example, let’s assume those are letters (indexed between 0 and 26 for all letters + a special end character) as in the great Andrej Karpathy “makemore<\/a>” series. <\/p>\n\n\n\n

Embeddings?<\/strong> For each element of that vocabulary, you’ve learned a representation of its (latent) characteristics, represented by a vector of size k. This vector is often called embeddings. Continuing with our example above, let’s consider an embedding of size 4 for each element (in our case, english letters) of the vocabulary, i.e. a tensor of dimension (27, 4)<\/p>\n\n\n\n

\"\"<\/a><\/figure>\n\n\n\n

Here is the gymnastic exercise<\/strong>: you have a toy batch containing 8 examples of size 3, where each number in the example are taken from vocabulary of size 27 . You also have an embedding matrix of dimension (27,4), where each raw is an embedding vector of size 4, for all of the 27 element of the vocabulary. For each element of the batch, you need to fetch its embedding vector, to end up with a batch which is a tensor of dimension (8,3,4) . This is illustrated below<\/p>\n\n\n\n

\"\"<\/a><\/figure>\n\n\n\n

Tensor indexing, the PyTorch way<\/h2>\n\n\n\n

Let’s first generate the two input tensors (the same as the two inputs on the left of the picture above ) :<\/p>\n\n\n\n

import torch\ntorch.manual_seed(18)\n# Create a random batch of shape (8,3) \n# with indexes between 0 and 26\nrandom_tensor = torch.randint(low=0, high=26, size=(8,3))\n# Create a random embedding matrix of shape (27,4): \n# one vector for each of the 27 indexes elements\nembeddings = torch.randn(size=(27, 4))<\/code><\/pre><\/div>\n\n\n\n
And now, let’s solve the gymnastic exercise. Take a deep breath, prepare the movement, and here you go:<\/p>\n\n\n\n
embedded_batch = embeddings[random_tensor]<\/code><\/pre><\/div>\n\n\n\n
Yes, that’s right. PyTorch allows to pass a full tensor as the index. And it works like magic.<\/p>\n\n\n\n
You can check the shape of the result, and observe it is indeed (8,3,4), as expected (see the picture above). Indeed, (8,3) is the shape of the initial batch, and for each element of it, we get the proper embedding vector of shape (1,4). <\/p>\n\n\n\n
Let’s validate that the first element of the result (embedded_batch[0,0] ) corresponds to the embedding vector of the index of the first element of the batch. This corresponds to this part of the picture:<\/p>\n\n\n\n
<\/p>\n\n\n\n
And sure enough, it worked \ud83c\udf89 :<\/p>\n\n\n\n
\"\"<\/a><\/figure>\n\n\n\n
What about TensorFlow?<\/h2>\n\n\n\n
In TensorFlow, it is of course possible to achieve the same result, but this is done a bit differently.<\/p>\n\n\n\n
The tf.gather function<\/h3>\n\n\n\n
Instead of injecting the batch directly as a (tensor) index in the embedding matrix, in TensorFlow we have to use a very powerful function:  tf.gather<\/a> .<\/p>\n\n\n\n
You can read the details of the documentation, but essentially, the equivalent of the following PyTorch indexing: <\/p>\n\n\n\n
embedded_batch = embeddings[random_tensor] <\/code><\/pre><\/div>\n\n\n\n
in TensorFlow would be: <\/p>\n\n\n\n
embedded_batch = tf.gather(embeddings,random_tensor)<\/code><\/pre><\/div>\n\n\n\n
And that’s all.<\/p>\n\n\n\n
Full equivalent TensorFlow code below :<\/p>\n\n\n\n
import tensorflow as tf\ntf.random.set_seed(18)\n# Create a random batch of shape (8,3) with indexes between 0 and 26\nrandom_tensor = tf.random.uniform(shape=(8,3), minval=0, maxval=26, dtype=tf.int32)\n# Create a random embedding matrix of shape (27,4): one vector for each of the 27 indexes elements\nembeddings = tf.random.uniform((27,4), dtype=tf.float32)\n# Solving the gymnastic exercise: creating an embedded batch with the tf.gather function\nembedded_batch = tf.gather(embeddings,random_tensor)\n# Validating the results\nprint(random_tensor)\nprint(embeddings)\nprint(embedded_batch.shape) # (8,3,4) which is the expected dimension\nprint(embedded_batch[0,0])<\/code><\/pre><\/div>\n\n\n\n
Hope you enjoyed the gymnastic lesson. Take some rest. Until the next one \ud83e\udd38 . <\/p>\n\n\n\n
References<\/h2>\n\n\n\n