How to set up a pymc environment on google cloud compute platform?

Hello @fonnesbeck.

In regards to google cloud platform. What do I have to do to use a GPU instance once spun up? Is there a special install of PyMC and Jax? Do I need to convert data to aesera tensors?

This is what I currently have in my instance.