fix: guard against unbounded tensor allocation in Layer.__call__

[HyperPS]

Summary

This PR adds a framework-level safeguard to prevent unbounded tensor allocations during layer execution, mitigating a potential denial-of-service vector when loading or executing untrusted .keras models.

Problem

Keras currently allows model-defined layers to request arbitrarily large output tensors during execution (e.g. via shape-controlled ops such as tf.image.resize). Because no framework-level limits are enforced, a malicious or untrusted model can trigger excessive memory allocation, leading to process crashes or resource exhaustion.

This is particularly relevant in deployment scenarios where .keras models are fetched from external sources (model hubs, CI pipelines, inference services) and executed automatically.

Solution

This change introduces:

• A configurable upper bound on tensor size via the KERAS_MAX_TENSOR_ELEMENTS environment variable (default: 1e8)
• A lightweight runtime validation step in Layer.__call__ that rejects outputs exceeding this limit
• Safe handling of dynamic shapes and symbolic tensors (no impact on tracing or graph construction)

Key details

• Enforcement occurs after eager execution and before layout transforms or activity regularization
• Symbolic tensors (KerasTensor) are explicitly excluded
• The default limit is conservative and can be tuned or disabled by users if needed

Security impact

• Prevents untrusted models from triggering uncontrolled memory allocation
• Mitigates CWE-770: Allocation of Resources Without Limits or Throttling
• Reduces denial-of-service risk in automated and multi-tenant ML environments

Backward compatibility

• No behavior change for normal models
• Dynamic shapes are unaffected
• Opt-out / tuning supported via environment variable

Related

• Security disclosure submitted via Huntr LINK
• CWE-770: Allocation of Resources Without Limits or Throttling

[gemini-code-assist]

Summary of Changes

Hello @HyperPS, 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 implements a critical security enhancement by introducing a framework-level mechanism to limit tensor memory allocation during Keras layer execution. This prevents malicious or untrusted models from consuming excessive memory, thereby mitigating denial-of-service risks and improving the robustness of Keras in deployment scenarios. The safeguard is configurable and intelligently handles dynamic and symbolic tensors.

Highlights

• Tensor Allocation Safeguard: Introduces a configurable safeguard to prevent unbounded tensor allocations during Keras layer execution.
• Runtime Validation: Adds runtime validation within Layer.__call__ to reject output tensors exceeding a defined element limit, configurable via the KERAS_MAX_TENSOR_ELEMENTS environment variable (default: 1e8 elements).
• Dynamic Shape and Symbolic Tensor Handling: Ensures that dynamic shapes and symbolic tensors (KerasTensor) are safely handled and explicitly excluded from the size validation, preventing impact on tracing or graph construction.
• Security Enhancement: Addresses a potential denial-of-service (DoS) vulnerability (CWE-770) by preventing untrusted models from triggering excessive memory allocation, enhancing security in automated and multi-tenant ML environments.

🧠 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.

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 introduces a crucial security safeguard against unbounded tensor allocation, which could lead to denial-of-service attacks. The implementation is solid, introducing a configurable limit via an environment variable and performing validation at an appropriate stage in the layer's __call__ method. My only suggestion is to improve the error message to be more informative for the user, in line with the Keras API design guidelines.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 51.72414% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.42%. Comparing base (c718248) to head (041ad70).

Files with missing lines	Patch %	Lines
keras/src/layers/layer.py	51.72%	10 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #22094      +/-   ##
==========================================
- Coverage   82.83%   78.42%   -4.41%     
==========================================
  Files         592      592              
  Lines       62991    63019      +28     
  Branches     9886     9895       +9     
==========================================
- Hits        52179    49425    -2754     
- Misses       8262    11142    +2880     
+ Partials     2550     2452      -98     

Flag	Coverage Δ
keras	78.29% <51.72%> (-4.37%)	⬇️
keras-jax	?
keras-numpy	56.62% <44.82%> (-0.01%)	⬇️
keras-openvino	37.47% <41.37%> (+<0.01%)	⬆️
keras-tensorflow	63.76% <51.72%> (-0.01%)	⬇️
keras-torch	62.53% <44.82%> (-0.01%)	⬇️

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]

There is more to memory allocation than what Layers can see (or what Keras can see for that matter).

Keras is the wrong software layer to apply this:

• Most computations allocate memory, and Keras may only knows about some of them. Only a compiler can have the full picture of the memory allocation. So this is at a lower layer than Keras.
• In terms of serving / inference. First, Keras is typically not how you perform inference. But if you do use it, you need to understand the exact characteristics of your model to know what the maximum input size for your model is. This is something that your inference service would enforce even before calling Keras, it doesn't belong in Keras.