fix(softmax): validate mask shape matches inputs shape

[rstar327]

Summary

• Add shape validation in Softmax.call() when mask is provided
• Previously, an incompatible mask shape would silently broadcast via numpy.where, producing wrong output shape
• Now raises ValueError with a clear message when shapes don't match
• Adds test for mask shape mismatch

Fixes #19722

[gemini-code-assist]

Summary of Changes

Hello @rstar327, 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 potential silent failure in the Softmax layer where an incorrectly shaped mask could lead to unexpected and erroneous output due to implicit broadcasting. By introducing explicit shape validation, the change ensures that such issues are caught early with clear error messages, improving the robustness and predictability of the Softmax layer's behavior when a mask is used.

Highlights

• Shape Validation: Added explicit shape validation within the Softmax.call() method to ensure that the provided mask has the same shape as the inputs tensor.
• Error Handling: Changed the behavior from silently broadcasting incompatible mask shapes to raising a ValueError with a descriptive message when a shape mismatch occurs.
• Test Coverage: Included a new test case to specifically verify that a ValueError is correctly raised when an incompatible mask shape is passed to the Softmax layer.

🧠 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/activations/softmax.py
  • Added a check in the call method to validate that the mask shape matches the inputs shape.
  • Introduced a ValueError to be raised if the mask and inputs shapes are not identical.
• keras/src/layers/activations/softmax_test.py
  • Added test_softmax_mask_shape_mismatch to confirm that a ValueError is raised when the mask shape is incompatible with the input shape.

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 adds a validation check in the Softmax layer to ensure that the mask has the same shape as the inputs. If the shapes do not match, a ValueError is raised. This prevents silent broadcasting and incorrect output shapes. A corresponding test case is also added to verify this behavior. The change is correct and well-tested. I have one suggestion to improve the readability of the error message to better align with the Keras API design guidelines.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 83.33333% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 82.94%. Comparing base (db4508d) to head (e834194).
⚠️ Report is 41 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/layers/activations/softmax.py	83.33%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #22295      +/-   ##
==========================================
- Coverage   82.94%   82.94%   -0.01%     
==========================================
  Files         594      595       +1     
  Lines       65361    66085     +724     
  Branches    10213    10316     +103     
==========================================
+ Hits        54211    54811     +600     
- Misses       8526     8658     +132     
+ Partials     2624     2616       -8     

Flag	Coverage Δ
keras	82.76% <83.33%> (-0.01%)	⬇️
keras-jax	60.81% <83.33%> (-0.43%)	⬇️
keras-numpy	55.01% <83.33%> (-0.41%)	⬇️
keras-openvino	49.10% <83.33%> (+10.21%)	⬆️
keras-tensorflow	62.04% <83.33%> (-0.45%)	⬇️
keras-torch	60.85% <83.33%> (-0.50%)	⬇️

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.

[hertschuh]

Thanks for working on this!

[rstar327]

Done, updated validation to allow broadcastable masks while disallowing masks larger than inputs.
Updated docstring and added tests for broadcastable, fewer-dim, mismatch, and too-many-dims cases.

[hertschuh]

Thanks!