PyMC3 Model Error When Using Einstein Sum

Questions
1

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[2], 1))
    u= pm.Normal("u", mu=0, sigma=1, shape=(1, X.shape[0]))
    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.

2

In

y_hat = gamma + np.einsum("jik,kl->ij", X, beta) + u

gamma and u are pymc distribution objects but np.einsum will return an ndarray. So to make y_hat an object in the pm.Model() context you should use

y_hat = pm.Deterministic("y_hat", gamma + np.einsum("jik,kl->ij", X, beta) + u)

However, if this doesn’t work (I haven’t checked), you may need to use a tensor algebra for the np.einsum calculation, such as this example.