Multiple TPUv2 devices in one training script?

Question

As part of the Tensorflow Research Cloud initiative, I have access to 100 TPU v2 machines with 8 TPUs on them (TPU v2-8s).

I need to achieve model data parallelism.