Fix one_hot symbolic output shape for IndexLookup

[Eruis2579]

Summary

Fixes #22336

StringLookup and IntegerLookup with output_mode="one_hot" were returning the wrong symbolic output shape for nested (higher-rank) inputs. Eager execution was correct, but building a model with keras.Input() produced (None, depth) instead of (None, d1, d2, ..., depth).

Changes

• keras/src/layers/preprocessing/index_lookup.py
  Updated compute_output_shape():

  • one_hot: keep all input dimensions and append vocabulary depth → input_shape + (depth,).
  • multi_hot, count, tf_idf: treat the last dimension as the sample dimension → input_shape[:-1] + (depth,).

• Tests

  • index_lookup_test.py: regression tests for one_hot symbolic shape and consistency with compute_output_shape / compute_output_spec.
  • string_lookup_test.py: regression test for StringLookup one_hot with nested input (issue repro).

[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]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical bug where StringLookup and IntegerLookup layers, when configured with output_mode="one_hot", generated an incorrect symbolic output shape for inputs with multiple non-batch dimensions. The changes ensure that the symbolic output shape accurately reflects the eager execution behavior by preserving all input dimensions and appending the vocabulary depth, thus resolving shape mismatches when building Keras models with keras.Input() and improving the reliability of these preprocessing layers.

Highlights

• Corrected Symbolic Output Shape: The symbolic output shape calculation for one_hot mode in IndexLookup (and its subclasses like StringLookup and IntegerLookup) has been corrected to properly handle higher-rank input tensors.
• Updated compute_output_shape Logic: The compute_output_shape method now correctly preserves all input dimensions and appends the vocabulary depth when output_mode is one_hot, ensuring consistency between symbolic and eager execution.
• Added Regression Tests: New regression tests have been added in index_lookup_test.py and string_lookup_test.py to validate the corrected symbolic output shapes for one_hot mode with nested inputs, specifically addressing issue StringLookup layer with output_mode="one_hot" produces incorrect symbolic output shape for nested input tensors #22336.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• keras/src/layers/preprocessing/index_lookup.py
  • Modified compute_output_shape to differentiate between one_hot output mode and other modes (multi_hot, count, tf_idf).
  • For one_hot mode, the output shape now appends the vocabulary depth to all input dimensions (e.g., input_shape + (depth,)).
  • For multi_hot, count, and tf_idf modes, the last input dimension is replaced by the vocabulary depth (e.g., input_shape[:-1] + (depth,)).
• keras/src/layers/preprocessing/index_lookup_test.py
  • Added test_one_hot_symbolic_output_shape_with_higher_rank_input to verify the correct symbolic output shape for IntegerLookup with one_hot mode and 3D input.
  • Included assertions to ensure symbolic and eager output shapes match (except for the batch dimension).
  • Added test_one_hot_compute_output_shape_multi_hot_consistency to confirm compute_output_shape behavior for both multi_hot and one_hot modes.
  • Added test_one_hot_compute_output_spec_preserves_input_dims to ensure compute_output_spec correctly preserves input dimensions for one_hot mode.
• keras/src/layers/preprocessing/string_lookup_test.py
  • Added test_one_hot_symbolic_output_shape_nested_input as a regression test for StringLookup with one_hot output mode and nested input, confirming the symbolic output shape matches eager execution.

Activity

• No human activity (comments, reviews) has been recorded on this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request fixes an issue with the symbolic output shape for IndexLookup layers when output_mode="one_hot" is used with nested inputs. The changes correctly preserve the input dimensions in the output shape. The added tests are good and cover the reported issue.

I've found a potential issue with an edge case that is not covered by the new tests. The logic in compute_output_shape for one_hot mode might be inconsistent with the runtime behavior implemented in encode_categorical_inputs when the last dimension of the input shape is 1. My review comment includes a suggestion to address this for better consistency.

[hertschuh]

I think the change is correct, however the tests are failing. Can you look into this?

[codecov-commenter]

Codecov Report

❌ Patch coverage is 66.66667% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.94%. Comparing base (b8909f8) to head (7a4021b).
⚠️ Report is 10 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/layers/preprocessing/index_lookup.py	66.66%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #22350      +/-   ##
==========================================
+ Coverage   76.76%   82.94%   +6.18%     
==========================================
  Files         595      595              
  Lines       65839    66084     +245     
  Branches    10277    10314      +37     
==========================================
+ Hits        50540    54813    +4273     
+ Misses      12879     8657    -4222     
- Partials     2420     2614     +194     

Flag	Coverage Δ
keras	82.77% <66.66%> (+6.17%)	⬆️
keras-jax	60.81% <66.66%> (-0.16%)	⬇️
keras-numpy	55.00% <33.33%> (-0.14%)	⬇️
keras-openvino	49.09% <0.00%> (?)
keras-tensorflow	62.04% <66.66%> (-0.14%)	⬇️
keras-torch	60.85% <66.66%> (-0.19%)	⬇️

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.

🚀 New features to boost your workflow:

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

[Eruis2579]

@hertschuh Please review my PR, All checks are passed

[Eruis2579]

@hertschuh , Could you check my updates

[Eruis2579]

@hertschuh , Please review my PR

[hertschuh]

Thanks for the fix!