Package: eindex
eindexConcept of multidimensional indexing for tensorsExample of K-means clusteringPlain numpydefkmeans(init_centers,X,n_iterations:int):n_clusters,n_dim=init_centers.shapen_onservations,n_dim=X.shapecenters=init_centers.copy()for_inrange(n_iterations):d=cdist(centers,X)clusters=np.argmin(d,axis=0)new_centers_sum=np.zeros_like(centers)indices_dim=np.arange(n_dim)[None,:]np.add.at(new_centers_sum,(clusters[:,None],indices_dim),X)cluster_counts=np.bincount(clusters,minlength=n_clusters)centers=new_centers_sum/cluster_counts[:,None]returncentersWith eindexdefkmeans_eindex(init_centers,X,n_iterations:int):centers=init_centersfor_inrange(n_iterations):d=cdist(centers,X)clusters=EX.argmin(d,'cluster i -> [cluster] i')centers=EX.scatter(X,clusters,'i c, [cluster] i -> cluster c',agg='mean',cluster=len(centers))returncentersTutorial notebookGoalsForm helpful 'language' to think about indexing and index-related operations. Tools shape mindsCover most common cases of multidimensional indexing that are hard to implement using the standard APIApproach should be applicable to most common tensor frameworks, autograd should work out-of-the-boxAim for readable and reliable codeAllow simple adoption in existing codebasesImplementation should base on fairly common tensor operations. No custom kernels allowed.Complexity should be visible: execution plan for every operation should form a static graph.
Structure of the graph depends on the pattern, but not on tensor arguments.Non-goals: there is no goal to develop 'the shortest notation' or 'the most advanced/comprehensive tool for indexing' or 'cover as many operations as possible' or 'completely replace default indexing'.ExamplesFollowtutorialfirst to learn about all operations provided.Click to unfold- how do I select a single embedding from every image in a batch?Let's say you have pairs of images and captions, and you want to take closest embedding:score=einsum(images_bhwc,sentences_btc,'b h w c, b token c -> b h w token')closest_index=argmax(score,'b h w token -> [h, w] b token')closest_emb=gather('b h w c, [h, w] b token -> b t c',images_bhwc,closest_index)To adjust this example for video not image, replace 'h w' to 'h w t'. Yes, that simple.- how to collect top-1 or top-3 predicted word for every position in audio/text?[most_likely_words]=argmax(prob_tbc,'t b w -> [w] t b')[top_words]=argsort(prob_tbc,'t b w -> [w] t b order')[...,-3:]- how to average embeddings over neighbors in a graph?# without batch (single graph)gather('vin c, [vin, vout] edge -> vout',embeddings,edges)# with batch (multile graphs)gather('b vin c, [b, vin, vout] edge -> b vout',embeddings,edges)- can eindex help with (complex) positional embeddings?If we're speaking about trainable abspos, it can be just saved asemb_hwcand added every time to a batch.
There is no need for indexing.But it can be very helpful for complex scenarios: for example in T5-relpos, when a bias is added to every logit before softmax-ing to compute attention?
That's simple to implement for 1d, andmuchharder for 2d/3d. Let's implement for 2d with eindex:N=Nonepos# [I, J] i jpos1=pos[:,:,:,N,N]pos2=pos[:,N,N,:,:]xy_diff=(pos1-pos2)%image_side# we make shifts positive by wrappingattention_bias=gather('i j head , [i, j] i1 j1 i2 j2 -> i1 j1 i2 j2 head',biases,xy_diff)Note that we indeed encounter relative position (shift in x and y), which is not done in most implementations that deal with flat sequence instead.In a similar way we could produce vector-shift attention (another typical version of relpos):vector_shift=gather('i j head c, [i, j] i1 j1 i2 j2 -> i1 j1 i2 j2 head c',biases,xy_diff)ImplementationRepo provides two implementation:array api standard. This implementation is based on astandardthat multiple frameworks pre-agreed to follow.
Implementation uses only API from standard, so all available operations support all frameworks that follow the standard.At some point this should become the one and the only implementation.Here is the catch: current support of array api standard is poor, that's why the second implementation existsnumpy implementationThis independent implementation works right now.Numpy implementation is great to test things out, and is handy for a number of non-DL applications as well.Development StatusAPI looks solid, but breaking changes are still possible, so lock the version in your projects (e.g.eindex==0.1.0)Related projectsOther projects you likely want to look at:tullioby Michael Abbott (@mcabbott) provides Julia macros with a high level of flexibility.
Resulting operations are then compiled.torchdimby Zachary DeVito (@zdevito) introduces "dimension objects", which in particular allow convenient multi-dim indexingeinindexis an einops-inspired prototype by Jonathan Malmaud (@malmaud) to develop multi-dim indexing notation.
(Also, that's why this package isn't calledeinindex)ContributingWe welcome the following contributions:next time you deal with multidimensional indexing, do this with eindexWorked? → great -let us know; didn't work or unclear how to implement → post indiscussionsif you feel you're already fluent in eindex, help othersguides/tutorials/video-guides are very welcome, and will be linkedIf you want to translate tutorial to other language and post it somewhere - welcomeDiscussionsUse github discussions for this projecthttps://github.com/arogozhnikov/eindex/discussions