I was planning to teach np.einsum to colleagues, by hoping to show how it would be reduced to multiplications and summations.
So, instead of numerical data, I thought to use alphabet chars. in the arrays.
Say, we have A (2X2) as [['a', 'b'], ['c', 'd']] and B (2X1) as [['e'], ['f']]
We could use einsum to create a matrix C, say like: np.einsum('ab , bc -> ac', A, B).
What I'd like to see is: it return the computation graph: something like: a*c + ..., etc.
Ofcourse, np.einsum expects numerical data and would give an error if given the above code to run.
tensordot has an example using strings for one of its arrays, taking advantage of the fact that 'a'*3 => 'aaa'. But einsum can't do anything with strings (tha'ts a compiled code issue).
Sometime ago I wrote a pure-python work alike, that parses the 'ij,jk->' string, and sets up the appropriate sum-of-products calculation. That includes extra debugging output. That might serve as a starting point for your task.
https://github.com/hpaulj/numpy-einsum
The latest einsum does some optimization, with some debugging help. np.einsum_path provides more information on that.