Implementation nlinalg.alloc_diag() in Aesara tensor


I have been trying to transcribe a model that I have in PyMC3 using the library Theano to PyMC4 using Aesara. However, I’m getting few errors, I think that my problem is to generate my covariance matrix. In PyMC3 I was using the following function

σ = pm.HalfNormal(‘σ’, sd = 10, shape = 4)
cov = tt.nlinalg.alloc_diag(σ)
μ = tt.stack([μ0, μ1, μ2, μ3]).T
joint_obs = pm.MvNormal(‘joint’, mu=μ, cov=cov, observed=ρ_hat)

However, Aesara has no nlinalg.alloc_diag(), is there any function or way that I could do the same?

Hi! If all you want to do is get a diagonal matrix, you could use broadcast multiplication together with an identity matrix, as in: cov = at.eye(4) * σ.

Alternatively, you could use np.diag_indices in conjunction with at.set_subtensor. at.set_subtensor will be more flexible if you need to make arbitrary covariance matrices, but it ends up being a bit longer in this case:

cov = at.zeros(4,4)
idx = np.diag_indices(4)
cov = at.set_subtensor(cov[idx], σ)
1 Like

Aesara also has at.fill_diagonal() which may help.


Thanks! Now, I can compile the model.

I didn’t know about that one! There’s also at.diag. Seems making diagonal matrices is popular enough to warrant several options.

1 Like