Softmax Regression with shared coefficients between classes

I appreciate the warning regarding shapes handling, but in my case that does not seem to be the issue. PyMC is able to infer the shapes correctly when I use my originally multidimensional Xs inputs.