I want to second the idea of broadcasting based on named axes/dimensions. I think it's a logical next step in the evolution of the array programming paradigm.
I particularly recommend checking out xarray. It has made my numpy-ish code like 90% shorter and it makes it trivial to juggle six+ dimensional arrays. If your data is on a grid (not shaped like a table/dataframe), I see no downsides to using xarray instead of bare numpy.
I've been writing my own low-level numeric routines lately, so I'm not up-to-date on the latest news, but there have been a few ideas floating around over the last few years about naming your axes and defining operations in terms of those names [0,1,2,3]. That sort of thing looks promising to me, and one of those projects might be a better conceptual fit for you.
[0] https://nlp.seas.harvard.edu/NamedTensor
[1] https://pypi.org/project/named-arrays/
[2] https://pytorch.org/docs/stable/named_tensor.html
[3] https://docs.xarray.dev/en/stable/