Simplex parameter for a Dirichlet distribution.
Real parameter for mean and sigma of Gaussian distribution.
If there are no non-reparameterizable random variables (as it seems to be in your case) you can use Trace_ELBO
. You can read more on this in SVI Part III: ELBO Gradient Estimators