Store tensors from TPU into variables, for use on CPU.
Store tensors from TPU into variables, for use on CPU.

New features:
* tpu_util.write_to_variable will make the tensor available on CPU.
* tpu_util.read_from_variable will retrieve the previously stored tensor.

PiperOrigin-RevId: 254087152

google-research/morph-net