Add Rematerialization to Keras

[divyashreepathihalli]

No description provided.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 49.45055% with 46 lines in your changes missing coverage. Please review.

Project coverage is 72.13%. Comparing base (90568da) to head (8c30ed7).

Files with missing lines	Patch %	Lines
keras/src/layers/layer.py	21.73%	35 Missing and 1 partial ⚠️
keras/src/ops/core.py	46.15%	7 Missing ⚠️
keras/api/_tf_keras/keras/ops/__init__.py	0.00%	1 Missing ⚠️
keras/src/backend/common/remat_scope.py	95.83%	0 Missing and 1 partial ⚠️
keras/src/backend/torch/core.py	50.00%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (90568da) and HEAD (8c30ed7). Click for more details.

HEAD has 4 uploads less than BASE

Flag	BASE (90568da)	HEAD (8c30ed7)
keras	5	3
keras-numpy	1	0
keras-torch	1	0

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #20743      +/-   ##
==========================================
- Coverage   82.01%   72.13%   -9.88%     
==========================================
  Files         557      558       +1     
  Lines       52016    52106      +90     
  Branches     8037     8056      +19     
==========================================
- Hits        42659    37587    -5072     
- Misses       7403    12671    +5268     
+ Partials     1954     1848     -106     

Flag	Coverage Δ
keras	72.05% <49.45%> (-9.78%)	⬇️
keras-jax	64.14% <48.35%> (-0.09%)	⬇️
keras-numpy	?
keras-openvino	29.91% <39.56%> (+0.01%)	⬆️
keras-tensorflow	64.76% <48.35%> (-0.03%)	⬇️
keras-torch	?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mattdangerw]

Might be worth implementing the higher level APIs in this PR too, as a way to validate that the low-level and high level can work together (and that this low-level API gives us what we want).

And maybe write a colab showing this functionality in action.

[fchollet]

Thanks for the PR!

[fchollet]

Add return section

[fchollet]

Also add it in the numpy and openvino backends, returning plain function results with no remat (for compatibility)

[innat]

@divyashreepathihalli
Thanks for the initial pull request for this. I noticed you're using the tf.recompute_grad API and attempt to provide high-level API for ease of use, which is great. A while back, I came across cases where people used various hacks with tf.recompute_grad for advanced modeling, like in EfficientDet; please check this comment.

To ensure broader usability, it might be worth considering such scenarios when designing the high-level API, so others don't have to rely on hacks, or at least minimize the need to do so.