pip install einindex==0.0.4



Write indexing operations in an elegant syntax inspired by Einstein notation.


Install with pip install einindex.

Then simply import the index function in a Python session:

from einindex import index


Given an array x with dimensions i and a set of indices idx with dimension j, x[idx] can be written as

index(x, idx, "i, [i]j->j")

Say we have an array x with dimensions (i, j) and indices idx with dimension j and are trying to compute the following y:

y = torch.empty(i)
for i_index in range(i):
    y[i_index] = x[i_index, idx[i_index]]

This becomes

index(x, idx, "i j, [j]i->i")

Limited support for multindexing is present.

Imagine if x is a batch of images with dimensions (batch, width, height) and we are trying to pick out one pixel from each image, so the indexing array has dimensions (batch, 2) and we want to compute y where

y = torch.empty(batch)
for batch_idx in range(batch):
    y[batch_idx] = x[batch_idx, idx[batch_idx, 0], idx[batch_idx, 1]]

This becomes

index(x, idx, "batch width height, [width height] batch->batch")