Allow tf.train.Checkpoint.write in tf.functions

We can't update the Checkpoint proto state (e.g. for tf.train.latest_checkpoint), but this at least throws an informative error and gives the user the option to write a checkpoint from a function without updating the Checkpoint proto.

Restore is more complicated, since we do restore-on-create. I think the right thing to do is to inherit the eagerness of the context where restore() was first called. Too much for this change.

PiperOrigin-RevId: 233082153