Pretty sure Numpy’s einsum[1] function allows all of this reasoning in vanilla numpy (albeit with a different interface that I assume this author likes less than theirs). Quite sure that first example of how annoying numpy can be could be written much simpler with einsum.
I think you could for linear solves into a "generalized Einstein notation", but the other option you have is to support more complex array types, in which case a batched linearsolve can be reframed as a single linearsolve by a block-diagonal matrix
Sure, but einsum needs a syntax and concepts of its own, and more importantly does not work if you need to do something else than a limited set of matrix operations.
It's what the "first example of how annoying numpy can be" does, using either np.linalg.solve in a loop or a cursed multidimensional index rearrangement.
The "cursed multidimensional index rearrangement" is mostly for doing inner product in the weird way demonstrated. Granted, it proves author's point - you need to be fluent in NumPy to write this and can't take the Go approach of "I refuse to learn anything you better make me up to speed in 5 minutes".
AiX = np.linalg.solve(A, X[:, np.newaxis, :])
Z = np.einsum("jk,ijk->ij", Y, AiX)
The NumPy 1.x code is very intuitive, you have A[i, j, n] and X[i, n] and you want to use the same X[i] for all A[i, j], so just add an axis in the middle. Broadcast, which the author very strongly refuse to understand, deals with all the messy parts.
Alternatively, if you hate `[:, np.newaxis, :]` syntax, you may do:
I, J, K, K = A.shape
AiX = np.linalg.solve(A, X.reshape(I, 1, K, 1)).squeeze()
Z = np.vecdot(Y, AiX)
I can read this faster than nested for-loops. YMMV.
https://dynomight.net/numpy/