Einindex

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

Usage

Install with pip install einindex.

Then simply import the index function in a Python session:

from einindex import index

Examples

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")