Fix #23110: GPU memory leak in torch tree utilities (closure captures tensor containers)

[utsab345]

Description

Fixes #23110.

Fixes a GPU memory leak caused by lambda closures capturing tensor containers in torchtree_impl.py and optree_impl.py.

Root cause

Both _dict_to_ordered_dict and traverse in torchtree_impl.py, and traverse in optree_impl.py use:

is_leaf=lambda x: x is not structure

The lambda captures structure (which holds GPU tensors) by reference. PyTorch's tree_flatten internally creates a self-referential helper closure that captures this is_leaf lambda, forming a reference cycle:

tree_flatten.helper → is_leaf lambda → structure (GPU tensors)
     ↑                                      /
     └──────── captures self ←──────────────┘

Python's reference counting cannot break cycles, so GPU tensors remain allocated until the cyclic GC runs, which is infrequent in tight inference loops.

Fix

Replace with:

structure_id = id(structure)
is_leaf=lambda x: id(x) != structure_id

The integer id() doesn't reference the container, so GPU tensors are excluded from the cycle and can be freed promptly.

Contributor Agreement

• I am a human, and not a bot.
• I will be responsible for responding to review comments in a timely manner.
• I will work with the maintainers to push this PR forward until submission.

[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 introduces changes to index_lookup.py to lazily initialize and dispose of mutable hash tables during adaptation, releasing TensorFlow resources. It also updates the tree traversal implementations in optree_impl.py and torchtree_impl.py to use object IDs for leaf checking. The review feedback highlights two critical issues in index_lookup.py: first, calling finalize_state() on an un-adapted layer with TF-IDF mode will raise an AttributeError due to uninitialized idf_weights; second, disposing of the hash tables will cause subsequent calls to update_state() to crash unless they are lazily recreated.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

I am having trouble creating individual review comments. Click here to see my feedback.

keras/src/layers/preprocessing/index_lookup.py (733-737)

If finalize_state() is called on an un-adapted layer (or when adapting empty data) with output_mode="tf_idf", self.idf_weights will not have been initialized. This causes an AttributeError when the code attempts to access self.idf_weights.value() inside the if block.

To prevent this crash, we should check if self.idf_weights has been initialized before accessing it, for example:

if self.output_mode == "tf_idf" and hasattr(self, "idf_weights"):
    self.idf_weights_const = self.idf_weights.value()

keras/src/layers/preprocessing/index_lookup.py (800-804)

Disposing of the mutable hash tables by setting self.token_counts = None and self.token_document_counts = None will cause subsequent direct calls to update_state() to crash with an AttributeError (since update_state() directly calls self.token_counts.insert(...) without ensuring the tables exist).

To make this robust, update_state() should be updated to call self._ensure_adapt_tables() at the beginning of the method so that the tables are lazily recreated if they have been disposed of.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.04%. Comparing base (eea8e6e) to head (635cbc6).

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #23206      +/-   ##
==========================================
+ Coverage   80.39%   84.04%   +3.65%     
==========================================
  Files         465      465              
  Lines       69188    69191       +3     
  Branches    11379    11379              
==========================================
+ Hits        55622    58153    +2531     
+ Misses      10742     8113    -2629     
- Partials     2824     2925     +101     

Flag	Coverage Δ
keras	83.87% <100.00%> (+3.66%)	⬆️
keras-cpu	83.87% <100.00%> (+4.40%)	⬆️
keras-gpu	?
keras-jax	57.92% <33.33%> (-0.35%)	⬇️
keras-numpy	53.70% <33.33%> (?)
keras-openvino	59.53% <33.33%> (-0.01%)	⬇️
keras-tensorflow	59.51% <33.33%> (-0.30%)	⬇️
keras-torch	58.67% <66.66%> (-0.40%)	⬇️
keras-tpu	?

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

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[DavidAriOstenfeldt]

See #23111