This is standard notation that's been used for decades.
For this I needed a good notation for functions applied to specific dimensions and broadcasting over the rest. Like softmax in a transformer.
The function chapter is still under development in the book though. So if you have any good references for how it's been done graphically in the past, that I might have missed, feel free to share them.
Tensors, as used in ML, are much closer to a key-value store with composite keys and scalar values, with most of the complexity coming from deciding how to filter on those composite keys.
Drop me a line if you're interested in a chat. This is something I've been thinking about for years now.