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