Einsum for Tensor Manipulation
swe-to-mle.pages.dev
swe-to-mle.pages.dev
It's almost like when we moved from writing machine code to structured programming: instead of being content to document that certain registers corresponding to certain semantically meaningful quantities at particular times during program execution, we manipulated only named variables and left compilers to figure out how to allocate these to registers? We're at that point now with array programming.
https://nlp.seas.harvard.edu/NamedTensor
I already replied to a couple of other comments with this, but you're exactly describing the Python library Xarray: https://docs.xarray.dev/en/stable/ I've found it to work quite well in practice, and I've gotten a lot of mileage out of it in my work.
But I do think it's fair to say they are in their infancy, and there is a missing theoretical framework to explain what is going on.
I anticipate name-free array programming will eventually be considered a historical curiosity for most purposes, and everyone will wonder how we put up without it for so long.
Overall, I agree with your larger point. I think this it's a natural progression and pretty analogous to assembly language leading to higher-level languages with a much better developer experience. In that case we went from fixed registers to developer-named variables, and now we can go from fixed numerically-indexed array axes to ones indexed by developer-chosen names.
res2 = dumbsum(query, key, 'batch seq_q d_model, batch seq_k d_model -> batch seq_q seq_k')Then it could have syntax highlighting, auto complete, and so on. The type system for such a language could possibly include verifying shapes at compile time.
What would a language comprised of .ein source files for manipulating tensors, which compiles down to the same low level ops, look like?
Here's einsum in 200 lines of Common Lisp. All einsum expressions are statically analyzed, checked for errors, and AOT compiled to machine code: https://github.com/quil-lang/magicl/blob/master/src/high-lev...
(einsum (A i j) (B i k) (C k j))
results in the the updates
A[i,j] = \\sum_k B[i,k]C[k,j],
which is equivalent to matrix multiplication.
@tensor A[i,j] = B[i,k] * C[k,j]
A matrix multiply written with this looks like this:
enum { i = 2, j = 0, k = 1 };
auto C = make_ein_sum<float, i, j>(ein<i, k>(A) * ein<k, j>(B));
This library is a lot less automated than most other einsum implementations. You have to explicitly control the loop ordering (in the example above, the `j` loop is innermost because it is loop 0). If you build a good loop order for your shapes, the compiler will probably autovectorize your inner loop, and you'll get pretty good performance. Control over the loop ordering is in general a useful tool, but it's probably a lot lower level than most users want.
* Succinct syntax and operators tailored to AI model definition
* Fully statically typed tensors, including generic functions over tensor dimension and batch dimensions (...)
* First-class hyper-parameters, model parameters and model arguments for explicit model specification
* Einops-style reshaping and reductions - tensor dimensions are explicit not implicit
Some of the einsum expressions I have seen for grouped multi headed/query attention is mind-boggling and they get shipped to prod.
scholar.google.com/scholar?hl=en&as_sdt=7%2C39&q=trace+"string+diagrams"+"Einstein+summation"&btnG=#d=gs_qabs&t=1714249856116&u=%23p%3Dbw0D4te4FGEJ
This article would have been better if the author had generated these kinds of diagrams, to show us what tensors the code was summing, and how. They could have easily used GPT to write the diagrammatic TiKz code, like they did for the fantasy style introduction.
But the diagrams would have illuminated the math, like real magic, while the fantasy writing didn't do any real magic because it only obfuscated the math. GPT does this when you don't give it the context to really explain things.
https://scholar.google.com/scholar?hl=en&as_sdt=7%2C39&q=tra...
Direct arXiv PDF link to first result:
Even better is if the library keeps track of the names of the tensor axes and the einsum is automatic.
[1] https://pytorch.org/docs/stable/generated/torch.einsum.html
[2] https://numpy.org/doc/stable/reference/generated/numpy.einsu...
One can easily dismiss use of this kind of stringly typed construct, but once you see how abstraction enables optimizations in implementation, it's kind of win-win.
Having used these (mostly translating code that used them) I see the power and benefit. I also find it takes a lot of mentally energy to get my head around them and makes readability harder.
Why so? They take what was implicit and make it explicit, so I would expect this notation to be far easier to understand than transposes and matmuls. You also get better variable names for axes, etc.
My complaint with Einsum variants in Python is that they are not generic enough -- they are specialised to multiplication & addition whereas generalizing slightly to allow other operators makes the code very expressive and elegant, while allowing you to still compile down to efficient code. For example, the Tullio library in Julia lang is absolutely a breath of fresh air.
What you do is name all the axes of an array, and then for all array computations, you address the axes by name. The great thing is that for the correctness of any of these operations, you never need to remember what numerical order the axes are in, since you just use the names. (For _performance_, the order can be important, in which case you can easily permute the axes to any order you want and then access the underlying NumPy/CuPy/whatever arrays if you need to.)