Add TPU testing

[chenmoneygithub]

No description provided.

[mattdangerw]

Thanks for this!

I think my main piece of feedback is that we should avoid consolidating all the TPU test out of tree with the models themselves. We should keep this testing right with the rest of the model testing, so we can keep it from getting stale (and make sure people adding new models add TPU tests).

From talking with you it sound like the main difficulty is that connecting to a TPU is slow and expensive, but I think we can work around this with pytest fixtures, which will allow you to do session level setup in a modular way. This is slightly awkward given that our actual tests unittest not pytest (see this doc).

But think I was able to get it working. I pushed a commit to this branch switching to tests directly in the model code using a session level tpu strategy. Take a look, modify as needed!

[mattdangerw]

lgtm. should we add tpu testing for other models, or leave that as a follow up?

[jbischof]

I like how the tests are now integrated with pytest flags! Would it add anything to do a couple of batches of model.fit for TPU compat?

[chenmoneygithub]

@jbischof Yes! I will make a followup PR to enrich our TPU testing's coverage.