freeze pymc models

Some pymc models give errors if they are not frozen when sampling with jax, because there are still dynamic shapes that are not supported by jax.

We should just freeze the model internally when we compile with jax with the pymc function

```
from pymc.model.transform.optimize import freeze_data_and_dims
```