I have the following model, where X is NxTxK and y is TxN:
with pm.Model() as model: gamma= pm.Normal("gamma", mu=0, sigma=100, shape=()) beta= pm.Normal("beta", mu=0, sigma=100, shape=(X.shape, 1)) u= pm.Normal("u", mu=0, sigma=1, shape=(1, X.shape)) r = pm.Gamma("r", alpha=9, beta=4, shape=()) y_hat = gamma + np.einsum("jik,kl->ij", X, beta) + u y_like = pm.Normal("y_like", mu=y_hat, sigma=r, observed=y)
However, when I run the model, I get the following error
ValueError: einstein sum subscripts string contains too many subscripts for operand 1
If I do the calculation
np.einsum("jik,kl->ij", X, beta) outside the pm.Model() environment, it works as expected so it seems that np.einsum does not work in the wrapper function.