1. Add several public apis that gets called at different phases of apply_gradients. These apis by default will not do anything, but users can override them for customized logic. 2. Make _compute_gradients() and _aggregate_gradients() public api.
1. Add several public apis that gets called at different phases of apply_gradients. These apis by default will not do anything, but users can override them for customized logic.
2. Make _compute_gradients() and _aggregate_gradients() public api.

PiperOrigin-RevId: 406470033

keras-team/keras,keras-team/keras