Repository navigation
perf(layers): fast path for single-InputSpec assert_input_compatibility - #23184
Conversation
|
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[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a fast-path optimization in assert_input_compatibility for the common case of a single InputSpec and a single tensor input. This optimization bypasses tree.flatten and the generic loop, reducing Python-level overhead during layer forward passes. Additionally, a comprehensive test suite has been added to verify that this fast path behaves identically to the general path across various input shapes, dimensions, and configurations. There are no review comments, and I have no feedback to provide.
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.
Codecov Reportโ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #23184 +/- ##
==========================================
- Coverage 85.33% 84.51% -0.83%
==========================================
Files 480 480
Lines 63874 63885 +11
Branches 10798 10802 +4
==========================================
- Hits 54510 53993 -517
- Misses 6697 7236 +539
+ Partials 2667 2656 -11
Flags with carried forward coverage won't be shown. Click here to find out more. โ View full report in Codecov by Harness. ๐ New features to boost your workflow:
|
fc292e1 to
d2f0a2b
Compare
d2f0a2b to
6d048aa
Compare
d9e5db3 to
81c6825
Compare
81c6825 to
f84a401
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refactors input compatibility checking in Keras by extracting a single-spec validation helper _check_input_spec and introducing a fast path for single-spec, single-tensor inputs to bypass flattening overhead. It also adds a comprehensive test suite in input_spec_test.py. The review feedback highlights critical edge cases in the validation logic: first, when allow_last_axis_squeeze is active, a rank mismatch can lead to silent truncation in zip and bypass shape checks; second, out-of-bounds axes in spec.axes can trigger unhandled IndexError exceptions. Suggestions are provided to explicitly validate rank equality and axis bounds, along with a corresponding test case.
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.
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Address the gemini-code-assist review on this PR and fix two latent validation bugs. Both predate this PR on master (the code was moved verbatim into `_check_input_spec`); fixing them here fixes the fast path and the general path uniformly. - Raise a `ValueError` instead of an unhandled `IndexError` when `spec.axes` refers to an axis that does not exist on the input. - Only squeeze one side, and only for rank mismatches of exactly one, when `allow_last_axis_squeeze=True`. The previous unconditional double-squeeze compared misaligned dimensions and silently accepted incompatible inputs such as shape `(2, 3, 99, 1)` against spec `(None, 3, 1)`, or `(2, 1)` against spec `(None, 3)`. Rank mismatches that squeezing cannot resolve intentionally remain non-raising here, diverging from the review suggestion: `Functional` (the only user of `allow_last_axis_squeeze=True`) validates ranks right after in `_adjust_input_rank` with a more informative error that includes the input path, a behavior pinned by `functional_test.py::test_rank_standardization_failure`. The suggested strict rank check would also reject documented-legal inputs such as shape `(1,)` against spec `(None, 1)`. Also align the new test file with keras test conventions (no module docstring, no section dividers, house naming) and expand regression coverage: out-of-bounds axes, squeeze dimension mismatches, and the Functional rank-mismatch error path. The new regression tests fail before the fix and pass after, on all four backends.
4be28f4 to
0072c18
Compare
Per review, so the diff reads as a modification of assert_input_compatibility plus an appended helper rather than a rewrite of the whole file, and shorten the fast-path comment to the suggested two lines.
The squeeze rebinds the local `shape`, so a rank N+1 input that failed the squeezed comparison was reported with the squeezed shape: `(2, 3, 99, 1)` against spec `(None, 3, 1)` said `found shape=(2, 3, 99)`, naming a shape the caller never used. Keep the original for the message. The comment above the zip claimed rank validation is left to `Functional._adjust_input_rank`. It is not: `assert_input_compatibility` runs earlier in `Layer.__call__`, so that error is unreachable for these inputs. State what the code actually does instead.
|
Thanks for the review @hertschuh! All requested changes have been applied in
|
|
Please fix the failing CI checks |
keras.src.backend no longer exports convert_to_tensor since keras-team#23642, so the new tests raised AttributeError on every backend.
fixed it |
JyotinderSingh
left a comment
JyotinderSingh
left a comment
There was a problem hiding this comment.
Please see the inline comments.
Additionally, please add a measurement to the PR description
The description does not say what the fast path saves. Please add a before and after measurement for assert_input_compatibility with a bare spec, and for a full Dense call. The saving on a full layer call is small, we need the number to decide if the new branch is worth it.
| spec_shape = spec.shape | ||
| # Report the shape the caller passed, not the squeezed one. | ||
| input_shape = shape | ||
| if spec.allow_last_axis_squeeze: |
There was a problem hiding this comment.
Please move the squeeze change out of this PR and into its own PR.
Functional.input_spec returns a list, so no Functional model reaches the new fast path. The fast path serves layers with a bare InputSpec. This change applies only to Functional models and to layers that set allow_last_axis_squeeze=True. The two changes do not depend on each other.
This is a breaking change. model(x) now raises for inputs that master accepts.
Examples: a (28, 28, 1) model called with a (4, 28, 28, 3) input, or a (3,) model called with a (2, 1) input. The new rule is correct. The PR never rejects an input that the docstring allows, and it never accepts an input that master rejects. But how many users it breaks is something we should investigate - which is why it should be in a separate PR.
Here, please keep master's squeeze logic (a verbatim move into _check_input_spec). The ValueError for an out-of-bounds axis can stay.
In the separate PR, add a Functional test for the same-rank case. No test calls a Functional model with an input of the same rank and the wrong last axis.
There was a problem hiding this comment.
#23920
Opened a new PR and fixed the other actionable.
The allow_last_axis_squeeze change is a behaviour change for Functional models and does not depend on the fast path, so it moves to its own PR. _check_input_spec is now master's per-input checks moved verbatim, plus the ValueError for an out-of-bounds axis. Add tests for the inputs that must leave the fast path and for a None input with a required spec, and tighten three assertions.
|
Thanks for addressing the reviews. The gain is small, on a full I'm okay either way. @hertschuh WDYT? |
To me, part of the value is the refactor extracting |
Conflict in keras/src/layers/input_spec.py: master moved the per-input checks into _check_input_spec (keras-team#23184). The shape check there now takes this branch's version, without the squeeze of both shapes.
Closes #23275
Summary
assert_input_compatibilityflattens both the spec and the inputs and runs a full input-count check before validating anything, even for the common case of oneInputSpecagainst one tensor. This adds a fast path for that case and extracts the per-spec validation into_check_input_spec, which the general path also uses._check_input_specis master's per-input checks, moved without change to the bottom of the file, so the diff reads as a change toassert_input_compatibilityplus a new helper.Behaviour change
One: an out-of-range
spec.axeskey now raisesValueErrornaming the axis, the input rank and the full shape. Previously it raised a bareIndexErrorfrom the tuple lookup.The
allow_last_axis_squeezechange from earlier revisions is no longer part of this PR. It changes which inputsFunctionalmodels accept and does not depend on the fast path, so it will be sent as a separate PR, as requested in review.Performance
assert_input_compatibility(spec, x, name)withspec = InputSpec(min_ndim=2, axes={-1: 64}), which is what a builtDensesets, and a(32, 64)tensor. Per call:The bare spec takes the fast path and saves 0.25 to 0.40 ยตs per call. A list of one spec stays on the general path, which now pays for the call into the shared helper: up to 0.09 ยตs.
A full
Dense(64)call on a(32, 64)input, which callsassert_input_compatibilityonce with a bare spec:Timed in place inside the
Densecall, the check takes about twice as long as in an isolated loop, and the fast path saves 0.2 to 0.5 ยตs there (2.65 ยตs to 2.14 ยตs on torch). That is 0.2% to 1.4% of the call, so the saving on a whole layer call is about 1%.Method: one process per build, master and this PR interleaved over 5 rounds, the best of 7
timeitrepeats per round, median over rounds. The last column is the difference between two master builds measured in the same runs; on the per-call table that difference is at most 2.3%. CPU rows are Python 3.11; the CUDA row is Python 3.14 on a shared GPU.Coverage
input_spec_test.pyadds 28 tests. Twenty-three callassert_input_compatibilitydirectly and cover the squeeze cases that master accepts, the out-of-bounds axis, aNoneaxis value, wrong axes, wrong shape, wrong dtype,min_ndim/ndim/max_ndim, non-tensor inputs, the optional-spec and empty-spec early returns, and that a bare spec and a one-element list spec behave identically, so the fast path and the general path agree. Two of them cover the guard on the fast path: nested inputs (a one-element list, a dict, two tensors) must take the general path, andNonewith a required spec must raise. The remaining five go throughDenseand throughFunctionalmodels.Contributor Agreement
Closes #23275above).