But how will this tool help? If you put the sample you used in a comment, then your confusion is wondering if that's in sync... and whether the input is like what you wrote.
IMO the sweet spot that's easy both for humans & computers is index notation. Instead of `tf.add(X, tf.expand_dims(Y, 1))` (their example, which I think also applies to 4-tensors in ways not trivial to see from the sample) something like "out_{r,c} = X_r + Y_c" is completely unambiguous, and quicker to type than making a sample.
I don't think you can do that in tf.einsum, nor einops. Is there a Python package which does this? In Julia you can write `@cast out[r,c] := X[r] + Y[c]` (with my package, but several others share this notation).