`pm.Model.compile_dlogp()` returns values in the same ordering as defined in `pm.Model()`

Questions v5
1

Hi there,

I have been doing something silly the past few days, and just worked out the fault and wanted to mention it here so that future me or someone wanting to use PyMC similarly to me can solve the same error.

I want to be able to define a pm.Model(), and use the compile_logp() and compile_dlogp() to help evaluate an objective function of a variational approximation algorithm that I am writing in JAX.

To reconcile what was wrong, I reduced the problem to learning a variational approximation for the mean \mu and precision \tau of a univariate Normal. I want to pass unconstrained parameters into the model, so I pass in \log\tau to the model, which is as follows:

with pm.Model(coords={"obs"=np.arange(len(data))}) as simple_gaussian:
     log_tau = pm.Normal("log_tau", mu=0., sigma=1.)
     tau = pm.Deterministic("tau", at.exp(log_p))

     mu = pm.Normal("mu", mu=0., sigma=1.)

     y = pm.Normal("y", mu=mu, tau=tau, observed=data, dims="obs")

I compile the derivative of the log probability function of this model: derivative_p = simple_gaussian.compile_dlogp().

To get the gradients concerning \mu and \log\tau, I call

derivative_logp({"mu": current_mu, "log_tau": current_log_tau}),

which in my mind would return \bigg (\frac{d\log p}{d\mu}, \frac{d\log p}{d\log\tau} \bigg ), but the ordering is actually the opposite. If I wanted the function to return the ordering I expected, I would need to redefine the model as

with pm.Model(coords={"obs"=np.arange(len(data))}) as simple_gaussian:
     mu = pm.Normal("mu", mu=0., sigma=1.)

     log_tau = pm.Normal("log_tau", mu=0., sigma=1.)
     tau = pm.Deterministic("tau", at.exp(log_p))

     y = pm.Normal("y", mu=mu, tau=tau, observed=data, dims="obs")

The compile_dlogp function returns the gradients in order of parameter creation in the model, and this ordering has nothing to do with the dictionary you pass in as the argument. This post is mostly a note for people (like myself) unfamiliar with using the compile_dlogp function.

If there is a way to get the values returned in the exact ordering as the appearance of the keys in the dictionary, that would be interesting!

Thanks,
Conor

2

You can specify the vars you want the gradient wrt to when calling compile_dlogp. If you specify all variables as a list you get the same function but with the output order specified in the list (I think)

1 Like
3

Thank you for your response: you are spot on! The ordering of the list given to the vars argument determines the order that grad_logp returns the gradients.

grad_logp = simple_gaussian.compile_dlogp(vars=[mu, log_tau])
grad_logp({"mu": current_mu, "log_tau": current_log_tau})
2 Likes