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:
GitHub - ctongfei/nexus: Experimental tensor-typed deep learning (also compile time type safety)
GitHub - ofnote/tsalib: Tensor Shape Annotation Library (numpy, tensorflow, pytorch, ...) (uses python type annotations).
Has anyone tried using named arrays with flux on GPU?
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.
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?
That looks very cool! I’ll dig into it some, thanks.