ericjang updated code to use pmap on multi-device TPU