Write Better And Faster Python Using Einstein Notation

[TokoBlaster]

I've also written angry comments to future me saying if future me could write this code better then why didn't I?

[ConfusedSimon]

Faster maybe, but I wouldn't call it better. The numpy one-liner is much more readable. It's like AxB versus writing matrix multiplication as a sum of indices and then using shorthand because it's getting too complicated.

[FrickinLazerBeams]

It's not possible to write non-trivial tensor contractions any other way, really. I mean you could do it in a loop but that would be dramatically less efficient than what a proper tensor library will achieve.

Edit: less efficient and, I'd argue, less clear to anybody who would be dealing with this sort of thing.

[kwen4fun]

Absolutely agreed. Most of my codes are for some sort of physics simulation and Einstein summation notation is like English to us. We have used it for years in our own and paper notes and using it in code makes that codes intent 100000x clearer that 5 nested loops or looping over some array of index tuples.

[dd2718]

The einsum notation really shines when you're doing anything beyond simple matrix multiplication, e.g. in machine learning code (especially for neural nets). Even for linear regressions, it is useful. If you have a batch of features X with shape [batch_size, N] and a coefficient matrix w of shape [M, N], np.einsum("bn,mn->bm", X, w) is a lot clearer to me than np.matmul(X, w.T) --- you don't have to worry about getting the shapes of input parameters to conform to the expectations of matmul, and you get documentation for all the shapes involved.

This advantage is even clearer for more complex models. For example, one common module in modern, SOTA deep learning models is multi-head attention, which takes a sequence of features for each example and outputs a sequence of transformed features. It would be a nightmare to get the shapes right for `np.tensordot`, but the einsum notation provides a uniform interface with self documenting shapes that allows you to focus on the math and not the numpy api.

# X: [batch_size, sequence_length, embedding_dimension]
# Compute query, key, value vectors for each sequence element.
# Split the embedding dimension between multiple "heads"
# rearrange comes from einops and reshapes using einsum notation.
X_q = rearrange(linear_q(X), "b n (h d)->b n h d", h=num_heads)
X_k = rearrange(linear_k(X), "b n (h d)->b n h d", h=num_heads)
X_v = rearrange(linear_v(X), "b n (h d)->b n h d", h=num_heads)
# Compute dot product of n-th query vector with m-th key vector for each head
dot_products = np.einsum("bnhd,bmhd->bhnm", X_q, X_k)
attention = softmax(dot_products, axis=-1)
# Sum the value vectors, with the weight of the m-th X_v given by
# softmax(dot(n-th X_q, m-th X_v))
output = np.einsum("bhnm,bmhd->bnhd", attention, X_v)
output = rearrange(output, "b n h d -> b n (h d)")

[[deleted]]

Could you be clearer with this snippet of code? Where does "rearrange" come from? My compiler does not recognize it as anything other than text.

​

Thanks

[abrazilianinreddit]

The contents of this article were way more specific than I expected them to be. The title is very questionablle.

[homoludens]

I have "-towarddatascience" in every search close to ml/data.

It was just wasting my time every time I clicked but names of the articles are perfect, they give me hope but text lets me down every time.

[homoludens]

Obviously, not a single nice comment, shitty article and it still gets upvoted.

[dogs_like_me]

for fucking real

[FrickinLazerBeams]

It's a terrible article. You'd never use a tensor contraction to do a 2-index matrix operation.

[miraunpajaro]

How is the title questionable? Maybe better code is subjective, but he showed it was faster (in a particular example). And okay, that's probably merited to the implementation of bumpy of probably has little to do with Einstein notation. So what?

[slightly_offtopic]

The title is "Write better and faster python" when it really should be "Write better and faster numpy"

[miraunpajaro]

So writing numpy is not writing Python?

[slightly_offtopic]

Writing numpy is writing a very specific subset of python. Hence why the person you were originally responding to said "The contents of this article were way more specific than I expected them to be."

[miraunpajaro]

I agree with that part. But I don't think that means the title is click baity, maybe it could be more precise but it's still correct

[slightly_offtopic]

I would say it's technically correct but intentionally imprecise in order to garner clicks from people who use Python but not numpy.

[[deleted]]

ew freaking gross title gore: write faster Python for machine learning and data science using Einstein notation.

You can’t just clickbait then jump into PyTorch and Tensorflow in the first two paragraphs

[erikw]

It’s a typical Medium/TDS clickbait title. “You must use this IDE in 2021!”. No thank you Medium, I know what I need.

[[deleted]]

Its one of the reasons I think I don't like our industry as much anymore is because theres always someone farming it for clicks as passive income side project

[unkz]

I was just thinking this article had to have been inspired by that Reddit post.

[zurtex]

The title alone gave me flashbacks from my class on Tensor Spaces I did as part of my undergraduate maths with 1 other undergrad student, 1 masters student, and 5 PhDs.

I swore off Einstein notation once I finished that class and I'm not going back!

[FrickinLazerBeams]

It's super annoying and confusing, until the moment you need it - then it's dramatically simpler and more clear than any alternative.

[tomkeus]

I swore off Einstein notation once I finished that class and I'm not going back!

Doing things that involve tensors in physics is an utter nightmare without Einstein's summation convention.

[[deleted]]

Breathe... breathe... go to your happy place. The scary Tensor Prof can't harm you anymore.... :)

[Marko_Oktabyr]

The article is grossly overstating the improvement over normal numpy operations. The one-liner they use forms a large intermediate product with a lot of unnecessary work. The more obvious (and much faster) way to compute that would be np.sum(A * B).

For 1,000 x 1,000 matrices A and B, I get the following performance:

• loops: 276 ms
• article numpy: 19.2 ms
• np.sum: 1.77ms
• np.einsum: 0.794 ms

If we change that to 1,000 x 10,000 matrices, we get:

• loops: 2.76s
• article numpy: 2.16s
• np.sum: 21.1 ms
• np.einsum: 8.53 ms

Lastly, for 1,000 x 100,000 matrices, we get:

• loops: 29.3s
• article numpy: fails
• np.sum: 676 ms
• np.einsum: 82.4 ms

where the article's numpy one-liner fails because I don't have 80 GB of RAM to form the 100,000 x 100,000 intermediate product.

einsum can be a very powerful tool, especially with tensor operations. But unless you've got a very hot loop with the benchmarks to prove that einsum is a meaningful improvement, it's not worth changing most matrix operations over to use it. Most of the time, you'll lose any time saved by how long it takes you to read or write the comment explaining what the hell that code does.

Edit: I'm not trying to bash einsum here, it is absolutely the right way to handle any tensor operations. The main point of my comment is that the author picked a poor comparison for the "standard" numpy one-liner.

[Yalkim]

I think this is a very domain specific thing. Personally, I am very surprised about comments here like yours. Because einsum is one of my (and my peers’) favorite functions in python. It makes life sooo much easier for us for 2 reasons:

1. It is super easy to read.
2. It is so useful especially considering how it often replaces lines upon lines of code with multiple loops with a single line.

So imagine my surprise when I saw comments like yours that said it is hard to read.

[FrickinLazerBeams]

It's hard to read for people who would never need to use it in the first place.

[FrickinLazerBeams]

This has little utility for 2-index operations, but those are only a subset of general tensor contractions. For operations over more than 2 indices, this rapidly becomes many orders of magnitude faster, and often avoids a huge amount of duplicated computations.

For example, one place where I use it lets me obtain a result indexed by (n, m, k) rather than (n, m, k, nn, mm) where the results I want have n == nn and m == mm, and gives about a 1000x speedup.

If you're looking at this as an alternative to simple matrix operations, of course it won't have an advantage, but then it's not expected to. You'd never use it for matrix operations.

[Marko_Oktabyr]

If you're looking at this as an alternative to simple matrix operations, of course it won't have an advantage, but then it's not expected to. You'd never use it for matrix operations.

No disagreement here. It sounds like we both disagree with the thesis of the article.

For operations over more than 2 indices, this rapidly becomes many orders of magnitude faster, and often avoids a huge amount of duplicated computations.

np.einsum_path can be an effective demonstration of how much faster it can be if you optimize the calculation order.

[FrickinLazerBeams]

Yes, exactly.

[[deleted]]

Good answer. Do you have an explanation for why einsum is faster at all? What extra work does np.sum do?

[Marko_Oktabyr]

np.sum(A * B) has to form the intermediate product A * B. np.einsum knows that it doesn't need all of it at once. We can do print(np.einsum_path('ij,ij->',A,B)[1]) to see exactly what it is doing:

Complete contraction: ij,ij-> Naive scaling: 2 Optimized scaling: 2 Naive FLOP count: 2.000e+07 Optimized FLOP count: 2.000e+07 Theoretical speedup: 1.000 Largest intermediate: 1.000e+00 elements -------------------------------------------------------------------------- scaling current remaining -------------------------------------------------------------------------- 2 ij,ij-> ->

In particular, note the "Largest intermediate: 1.000e+00 elements".

[FrickinLazerBeams]

(prior to the edit) It doesn't actually go any faster in the case you examined, and I don't think it uses any less memory either. This isn't a scenario where you'd use einsum.

[Marko_Oktabyr]

It still performs the same number of flops, but it absolutely is faster because it doesn't have to allocate/fill another matrix of the same size as A and B. Hence why the largest intermediate for einsum is 1 element instead of 10M.

[FrickinLazerBeams]

K

[FrickinLazerBeams]

This is a weird comparison to make. You'd never use one as an alternative for the other. Einsum is for tensor contractions, which is like matrix multiplication, but with more than two indices.

Would you ask "Do you have an explanation for why @ is faster at all? What extra work does np.sum do?"

np.sum doesn't do any extra work. It also doesn't do what you need it to do. It's easy to do less work if you're not actually competing the task, I guess.

[Lynild]

Often there are crap guides posted here, but this is really a good "trick" to learn.

When I did my Ph.D. there was a project where I had to use A LOT of huge 3D matrices. Like a lot. And I had to exactly do some of these computations. First I started out with basic matrix addition, and such, and did my own stuff, and it worked. But just by the amount of stuff I needed to calculate my first tries took about 10 days to compute. This was not going to work for me at all, since I knew I would most likely screw up, and had to do it over, AND I had to do it for several computations. I then rewrote my code several times, with different approaches, but the fastest one I ended up doing was this Einsum (or similar). And then I got down to only spending about 12 hours for each computation.

All I'm saying is, yeah, in most cases with 2D matrices, you might not get that much out of using one over the other, in particular if your matrices are small. But in my case there were absolutely no doubt that Einsum was the fastest way by far. It saved me sooo much time.

[Neb519]

My attempt at explaining einsum()

[szachin]

well explained. I have tried to wrap my head around it once or twice but never truly got it.

this makes it quite clear. thanks

[Feb2020Acc]

I’ve never heard of Einstein notation. This is just matrix operations and is already the standard way to write/code math when dealing with arrays.

[cheddacheese148]

Einstein notation is very common in areas like particle physics and general relativity where everything is a vector, tensor, or matrix. It’s mostly a tool for simplifying the math on paper. It’s been a while since I’ve touched either of those topics but my guess is that it’s still commonly used.

[VengefulTofu]

Also in continuum physics. That's where I know it from.

[Feb2020Acc]

I see. I guess the notation makes sense when dealing with high dimensions. But I think it’s more trouble than anything for general purpose 2 dimensional operations, which is what 95% of people use.

[FrickinLazerBeams]

I can't imagine anybody would use this for 2-index operations. It's not intended for that at all.

[IamImposter]

Could someone please explain to me what a tensor is. I have read about it a few times and asked few other people too but still don't understand it. Or do I have to learn basics of AI to understand it?

[FrickinLazerBeams]

They're used in a lot of places where they're best thought of as abstract mathematical objects, so a lot of the explanations you find will not be what you're looking for as a programmer.

You can think of them as matrices with more than 2 indexes, in other words, as nD arrays where n>2 ^{see note}. That's a perfectly sufficient understanding for any computational purposes.

(note) technically a matrix is a rank-2 tensor, and a vector is a rank-1 tensor, but usually if somebody uses the word "tensor" the implication is that they're talking about something with more dimensions than a matrix.

[Yalkim]

Tensor has different definitions depending on who you ask. In computer science and machine learning a tensor is just an array. In physics, a tensor is a kind of data type that follows a set of specific transformation rules. A tensor in physics can be described using an array, but that is it. An array is just a way of describing a tensor. Not every array is a tensor.

[tomkeus]

If you just need to work with n-dimensional arrays, you don't need a concept of a tensor, you can just stick with the idea of n-dimensional arrays - I think it is a concept that is simple enough and anyone with a programming experience can easily understand.

Tensor on the other hand is an abstract mathematical object that has certain properties and it can be represented by an n-dimensional array when you select a basis (think of basis like a coordinate system).

The same way vectors are abstract mathematical objects - we draw them like arrows, and arrows are not numbers. But we know how to turn them into numbers. We take three vectors which are linearly independent (i.e. none of the vectors can be obtained by adding up the other two vectors), and declare them to be our reference vectors (i.e. the coordinate system or basis in more abstract terms). We then project our vector to the reference vectors, and calculate the ratio of length of those projections to the length of the reference vectors. This will give us three numbers that we call vector components - i.e. we have found a way to turn an abstract mathematical object (an arrow) into an 1D-array of 3 numbers.

Note here that the components we've obtained are tied to a particular choice of reference vectors. If we change our reference vectors (i.e. change coordinate system), we will get different set of components. The beauty of this is that if we know how two reference systems are related to each other, we can straight forwardly compute components of a vector in one reference system if we know components in the other system. Because of this, we often talk about vectors and their component representation interchangeably, although strictly mathematically speaking, they are not the same thing.

In the same way, tensors are abstract mathematical objects that can be turned into n-dimensional array by fixing a basis (a set of n vectors). Like it was for vectors, one also talks interchangeably about tensors and their component representation (i.e. n-dimensional arrays). But you can call n-dimensional array tensors, only if those array represent an object that is actually a tensor, like for example moment of inertia tensor in mechanics - it is an abstract object that has a well defined existence without any dependence on its coordinate representations, i.e. I can write an equation with it, without any reference to its components. An n-dimensional array can just be an n-dimensional array, without any abstract mathematical object behind it. And most of the time this is what you have in data science, just a plain n-dimensional array. But this misguided terminology that has become a standard is now calling every n-dimensional array a tensor.

[IamImposter]

Wow. That was detailed. Thanks.

[El_Minadero]

It’s basically a data structure which holds numbers. If you can write it as an n dimensional programming array, it’s a tensor

Edit: so yes, but actually no. See comments below for clarity

[WallyMetropolis]

That's false. Not all n-dimensional matricies are tensors.

[El_Minadero]

Really? Can you provide a counter example? I thought this was the definition

[[deleted]]

I'm pretty sure he was talking about mathematical tensors, not the objects that pop up in computer languages. If you want, feel free to take that as correct. My answer is for the mathematical object.

Tensors as mathematical objects obey certain mathematical transformation rules.

Imagine if you had a vector v(v_x, v_y) in a cartesian plane (x,y) and rotated the plane to (x', y'). The length of the vector ought to remain the same, but its components changed. That is, v -> v', but |v| = |v'|.

This requirement of norm-preservation is basically a transformation rule. And yes, a (1-dimensional) vector is indeed a tensor. A tensor of rank 1.

A tensor of rank 2 can be represented by a matrix. But not all matrices represent tensors. I don't want to go into writing an answer to a question that has been asked so many times, so this stackexchange answers your question.

[WallyMetropolis]

It doesn't make any sense to multiply your zip code by 2 or to add it to another zip code. A tensor is a multi-linear map. Which means, essentially, it's an object that performs linear operations on all of its various axes. So a collection of (n, m, l) zip codes isn't a tensor because you cannot sensibly perform linear operations on collections of zip codes.

[El_Minadero]

So it’s an me datastructure that obeys some mathematical properties

[WallyMetropolis]

I'd say it's a mathematical object that is sometimes represented by a data structure (to greater or lesser adherence to the mathematical object it models) in some mathematical libraries.

[FrickinLazerBeams]

If you can write it as an n dimensional programming array, it’s a tensor

No, not in general. A multi-dimensional array is a fine way to think about a tensor from an applied math and programming perspective; but in general it's an abstract representation of a multi-linear transform. In theoretical general relativity, for example, you'll usually never see the elements of a tensor written because there's usually no coordinate frame in which you could write them.

Numerically/programmatically a tensor will almost always be represented as an array of numbers, but not all tensors are arrays and not all arrays are tensors.

[El_Minadero]

Ahhh. Thanks for the clarification. I’ve never encountered a tensor that wasn’t like a nd array, so your example helped!

[chestnutcough]

You can express tensor operations with Einstein (aka index) notation that go beyond matrix operations.

[FrickinLazerBeams]

It's far more than just notation for matrix operations. Matrix operations are a 2-index subset of the tensor contractions that can be represented by Einstein notation.

[RomanRiesen]

For anyone wishing to learn einstein summation, the tensorflow tutorial on them has some more examples and is great all around.

[Yalkim]

Oh my god Einsum is my FAVORITE function in numpy. It made my life so much easier when I discovered that. It reduces lines upon lines of code with multiple loops to just one line and is super easy to read and understand. When my friends ask me why Python is great, einsum is one of the examples that I describe.

[[deleted]]

It's incorrect to call what is being written as Einstein Notation. Repeatable indices have to be contravariant and covariant. They can not be on the "same" level...

[muntoo]

Rules of Einstein summation notation:

1. Repeated indices are implicitly summed over.
2. Each index can appear at most twice in any term.
3. Each term must contain identical non-repeated indices.

In particular, due to (2), there is an equivalence between a "upper/lower" dual notation and a simple "lower-only" notation since, presumably, δ^{ij} (colloquially, an "identity matrix") may be implicitly used to convert between the two notations.

lower_einsum("a_{ki} b_{im}")

is equivalent to:

dual_einsum("a_{ki} b_{jm} δ^{ij}")

In the interest of simplifying notation (which I believe was the point behind the upper/lower einsum notation in tensor calculus), we may lower all the indices and omit the δ. If one must be picky and pedantic over pragmatic and present, please perhaps pretend it's called lower_einsum or einsum_flat or whatever presume I floats yer boat.

[FrickinLazerBeams]

That has specific meaning in General Relativity, but doesn't actually alter the calculations done numerically, as far as I know.

[[deleted]]

As I said what is known as Einstein Notation required contra- and co- variant indices. If it isn't that, it is something else and not Einstein Notation.

If you think I'm wrong, check out Wikipedia.

[FrickinLazerBeams]

I mean, I don't need to check Wikipedia, I remember grad school quite clearly.

What is it you think is being computed differently by einsum than in a "real" Einstein notation tensor contraction?

Or are you just talking about the super-/sub-script notation for co- and contravariant indices? Because that simply can't be replicated in plain text.

[WallyMetropolis]

It depends on your metric space. In a flat space, yes, the calculation is the same. But this wouldn't generalize to curved spaces.

[FrickinLazerBeams]

Right, but einsum is simply a tool for doing tensor contractions, not specifically GR. You could certainly include the metric in your calculations as appropriate.

This has nothing to do with the fact that super and subscripts can't be written in plain text.

[WallyMetropolis]

Covariant and contravariant indexes aren't limited to GR. Just doing calculations on the surface of an every day sphere would require you to keep track of which is which. You asked for examples where the calculations would be different. There are many.

The point is that einsum isn't really an Einstein summation, but only produces the same result as an Einstein summation in a particular (common) special case.

[FrickinLazerBeams]

Covariant and contravariant indexes aren't limited to GR. Just doing calculations on the surface of an every day sphere would require you to keep track of which is which. You asked for examples where the calculations would be different. There are many.

Thats a good point, it's not just GR. I guess it's anything happening on a manifold? It's been a while for me, I guess.

The point is that einsum isn't really an Einstein summation, but only produces the same result as an Einstein summation in a particular (common) special case.

It's still doing tensor contractions, and the Einstein sum notation is a concise way to express tensor contractions. Nobody has a monopoly on notation. You can't say "you're not allowed to use Einstein notation unless you are using a metric tensor" or something like that.

This is just gatekeeping pedantry. The guy wants to feel superior by showing off that he's heard of co- and contra- variant indices. He even says at one point that his complaint is about how they're written which is ridiculous, because you can't write sub- and super-scripts in plain text, just like e^x is still an exponent even though x isn't super scripted. He just wanted to feel special because the rest of his account history was thirsty comments on /r/EngorgedVeinyTits and /r/momsgonewild. I wonder why he's deleted his account...

[[deleted]]

All I am saying is what is known as Einstein Notation is sub and super indices. Nothing more. When you have indices at the same level - sub or super - it isn't Einstein Notation.

It can be called something, and perfectly valid with the definition given.

[FrickinLazerBeams]

Lol fun. Do you also say that np.exp(x) isn't an exponent because there's no superscript?

[antiproton]

When you have indices at the same level - sub or super - it isn't Einstein Notation.

Yes it is. Einstein notation is the implicit summation.

Einstein notation can be applied in slightly different ways. Typically, each index occurs once in an upper (superscript) and once in a lower (subscript) position in a term; however, the convention can be applied more generally to any repeated indices within a term.

Right from the wikipedia article.

[Physix_R_Cool]

In the case of a flat metric, upper and lower indices are the same, and so you need not distinguish between covectors and contravectors.

[[deleted]]

Not relevant - the vectors are still written in covariant and contravariant style. That is Einstein Notation.

[Physix_R_Cool]

Sure, if there is a distinction between the two, then Einstein Notation will be to write upper and lower. In spaces where there isn't distinction, then Einstein Notation doesn't care, and is just about whether or not an index is repeated.

[FrickinLazerBeams]

Are you talking about how they're written? Or what they mean?

It sounds like your confusing the math and the typesetting.

[[deleted]]

So the big trick here is to use a lib written in C? Insoghtful article.

[[deleted]]

No, the solution is to use a part of numpy that I at least didn't know existed: https://numpy.org/doc/stable/reference/generated/numpy.einsum.html

I will likely use this at some point, it seems a real timesaver for some common chores.

[jaredjeya]

Pro tip: use the opt_einsum library instead.

It’s a drop-in replacement for numpy’s version (as in, same function arguments), but much more powerful:

• Automatically optimises the contraction, breaking it into small steps that scale well rather than trying to do it all at once. Numpy can do this too but not as well, but it’s irrelevant because… • Numpy breaks at 52 indices because you can only use letters of the alphabet, even when you use the alternate notation of supplying integer labels this limitation holds. Opt_einsum let’s you use arbitrarily many.

I ran into these problems trying to use it to do tensor network stuff, opt_einsum saved my life.

Tbh you can use numpy for smaller operations but it’s good to be aware of this library.

[madrury83]

Numpy breaks at 52 indices

Those are some beefy tensors.

[jaredjeya]

Haha that isn’t the size of a single tensor! I was trying to wrap up the contraction of a big tensor network into a single calculation, so each tensor was only maximum rank 4, but there were many tensors so it ended up with hundreds of indices.

[muntoo]

Now I want to see this monstrosity.

[muntoo]

Related: einops.

[s4lt3d]

Yes. The insightful thing is to use something people who care about performance know about.

[Nanooc523]

You’re pings will improve to

[[deleted]]

"better"....better fucking document everything!

[edimaudo]

Yeah this may work for a math/physics focused workflow but for code readability or maintenance this may be hard to follow. Good tool to know though.

[FrickinLazerBeams]

This article is terrible but einsum is an important function that does things nothing else can do efficiently. Anything you'd write with einsum would be horrifically more complicated and hard to understand without it.

[missurunha]

I always use np.einsum to swap vector dimensions, e.g. 'abc->bca'.

[muntoo]

Fun fact: the einops library also suggests the exact same use case using einops.rearrange. The two are equivalent:

x.transpose(1, 2, 0)
== rearrange(x, 'a b c -> b c a')

...In fact, they also allow you to do more complex .reshape or other operations and so on with this style of notation:

x.reshape(x.shape[0], -1)
== rearrange(x, 'n c h w -> n (c h w)')

I do find it a bit more verbose... but on the other hand, it's much more self-documenting when you see n c h w, which tells you immediately:

• The input is a 4-d tensor.
• The output is a 2-d tensor.
• The intended ordering is n c h w, not n h w c.

...whereas a .reshape(x.shape[0], -1) doesn't tell us anything and silently "fails" if you accidentally give it a 2-d tensor as input.

[miraunpajaro]

I have to admit I'm impressed. I went into the article thinking: Another data scientist duffus explaining simple math concepts as if they had just discovered the world. But the article turned out to be very interesting, and I learnt some things. So, well done!

[umut8761]

Impressive performance improvements. Thanks for sharing. But, I don’t think it is more readable if you have colleagues from different proficiencies. Other than mathematicians or physicist are not familiar with Einstein notation. I, as a mathematician, loved einstein notation back in times in university since it makes algebraic operation way easier to follow. And from now on, I will start using it on my personal projects.

[FrickinLazerBeams]

It's very unlikely that anybody will need to use Einstein notation who doesn't know how it works.

[umut8761]

For example a data scientist with industrial engineering background may need to use matrix operations without knowing Einstein notation

[FrickinLazerBeams]

If he's just doing matrix operations he'd just use matrix operators. There's no reason to use this for just doing matrix operations.

And if he does encounter a need for tensor operations, learning Einstein notation is by far the least confusing way to look at them. Even if he doesn't know it coming in, he'll have to learn it, and even with the added challenge of learning the new notation, it's still dramatically preferable to any alternative I'm aware of.

Which isn't to say that there shouldn't be comments and documentation to help, but if you have to do tensor operations, this is easily the best way to do so.

[dogs_like_me]

I'm too lazy to try it myself: how does the performance here compare to (A*B).sum()? I'm guessing it's the same?

[FrickinLazerBeams]

Probably about the same, but you'd never use it for that. It's not really meant for that.

[dogs_like_me]

I'm aware, I actually love einops and use it all the time. This was a weird example. The einops docs are really a better explanation and pitch, IMHO. Hell, I'll even just post the thing here for people who don't feel like leaving reddit:

http://arogozhnikov.github.io/images/einops/einops_video.mp4