Fix memory leak by purging orphaned CustomGradients in clear_session

[opabhijeet]

This PR fixes the memory leak that happens when training Keras models in a loop (e.g., repeatedly calling model.fit()).

After doing some digging into the memory graph, I tracked down exactly why clear_session() wasn't freeing up the memory. The issue stems from how AutoGraph handles custom gradients (like CustomGradient-42358). It registers these closures permanently in the global tensorflow.python.framework.ops._gradient_registry. The problem is that these closures capture SymbolicTensor objects, which hold the Keras FuncGraph and its PyBind11 C++ wrappers.

Because Python's garbage collector can't traverse PyBind11 objects, that single global registry entry ends up pinning the entire dead execution graph in memory forever, even after Keras destroys the actual Model.

The Fix:
This PR just updates clear_session() to explicitly drop these orphaned custom gradients from the global registry. Once that anchor is severed, the Python GC immediately tears down the FuncGraph and frees the C++ memory.

(Note for reviewers: I opted to purge all dynamically generated CustomGradient entries since clear_session() acts as a global reset. If you'd prefer to filter these more strictly by inspecting the closure graphs, just let me know and I'm happy to update it!)

Related Issues

Fixes tensorflow/tensorflow#80895
Fixes #286

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request updates the clear_session function in tf_keras/backend.py to remove dynamically generated custom gradients from the global registry, which helps prevent memory leaks caused by circular references. A review comment identifies a potential RuntimeError when iterating over the gradient dictionary and suggests using a list snapshot of the keys to ensure thread safety.