You're probably better off using Google Cloud Machine Learning, which might by running on the custom TPU accelerator (not sure).
I realize that beta means beta, but there's very little working code out there that runs properly (especially Google's sample code), and in many cases their documentation is either outdated or just plain missing.
It just doesn't feel like it's baked, even for one of their betas.