Type safe and expressive "tensor" code


#1

There’s been some buzz lately about more expressive, and type safe array code for deep learning. One specific initiative has been the use of arrays with named dimensions:

http://nlp.seas.harvard.edu/NamedTensor
https://github.com/ctongfei/nexus (also compile time type safety)
https://github.com/ofnote/tsalib (uses python type annotations).

Has anyone tried using named arrays with flux on GPU?


#2

I have not, but as a budding macrologist I have another idea: You could extend the various index-notation macros to check at parse-time (if that’s the right term) that objects defined with certain indices are used again that way.

Papers with lots of many-index tensors often adopt some convention such as μ,ν=0...3 (spacetime directions, say) vs. M,N=0...9 vs. a,b,c = 1,2,3 (numbering some objects). Then if you define g_{μν}=... and later refer to g_{ab} this sets off a giant warning sign in referee #2’s brain.

We could have a similar warning if we made a version of @einsum etc. which, when you write @einsum_checked g[μ,ν] := something[i,μ,i,ν], saves the names of g's indices to some dictionary. And then complains if the next usage is different, by more than 2 places in the alphabet, or something.


#3

For your possible amusement: icheck.jl

This provides both the “referee #2” check mentioned, and (optionally) inserts run-time checks that sizes (along axes named with a particular letter) are as they were earlier. Very much the first draft, but perhaps not so far from what “tsalib” offers?


#4

That looks very cool! I’ll dig into it some, thanks.