[TF2] Remove training.partially_restore_from_checkpoint

[boomanaiden154]

Migrate training.partially_restore_from_checkpoint to the new TF APIs. It may or may not work currently as it has not been thoroughly tested and the new APIs natively support most of the functionality which would mean we can remove that function.