Checkpointable: Handle Optimizer non-slot variables

Overrides the Checkpointable dependency-gathering logic to key Optimizer dependencies to the current graph.

Moves my Checkpointable Adam prototype out of contrib.

Right now there is no check that loading all happens in the same graph. This would be easy enough to do (save a Graph ID with the _Checkpoint object), but it's not clear to me that it's useful; doing deferred restoration in whichever graph the variable is created in seems reasonable. (Let me know if you disagree)

PiperOrigin-RevId: 187378372