System information

• TensorFlow version (you are using): 2.1
• Are you willing to contribute it (Yes/No): Yes

Describe the feature and the current behavior/state.
I have implemented a version of gradient check pointing for TF keras sequential models (with future plans to extend it for Keras functional API and custom models). The PR can be found here - tensorflow/addons#1600. I initially envisioned it as a package in TF addons repo but reviewers felt it was not right place and could potentially go in as a fix for the existing recompute_grad functionality in TF core -

Here are the issues with the existing implementation of recompute_grad in TF core and my solutions to those

1. Issue - Not usable. Most people have no idea how to use it. No docs or tutorials that explains how to use it.
   Solution - My PR provides a notebook tutorial to demonstrate how to use the implemented functionality.
2. For people who did figure out how to use it, no memory savings was observed.
   Solution - My PR provides links to results with observed memory savings. Caveat - only CPU profiled results available. GPU and TPU results need to be done.
3. There is probably an expectation that the user explicitly has to partition the model and decorate each partition. This is not user friendly and can make tasks such as transfer learning difficult.
   Solution - My PR expects no explicit partitioning of the model. The user just needs to add a single decorator to the model. That is it.
4. No checkpointing functionality is implemented.
   Solution - My PR implements the checkpointing functionality that allows the user to balance the tradeoff between memory and compute time.

Does it make sense to port the PR to TF core?

Will this change the current api? How?
The existing implementation can potentially be shoe horned into the existing API for recompute_grad if desirable.

Who will benefit with this feature?
Anyone who wants to train models in resource constrained environments.

Any Other info.