Repository navigation
perf(backend): O(1) ALLOWED_DTYPES membership and a fast path for canonical dtype strings - #23182
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 optimizes the standardize_dtype function by converting ALLOWED_DTYPES to a frozenset for O(1) lookups, introducing fast paths for common dtype formats, and implementing a module-level cache (_DTYPE_CACHE) for resolved non-string dtypes. However, the new string fast path introduces a regression where framework-prefixed dtype strings (such as "torch.float32" or "jax.numpy.float32") are no longer correctly resolved and instead raise a ValueError. The feedback suggests updating the string fast path to handle these prefixed strings and adding corresponding test cases to prevent future regressions.
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.
0e93741 to
1343b06
Compare
Codecov Reportโ
All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #23182 +/- ##
==========================================
- Coverage 86.88% 84.31% -2.58%
==========================================
Files 477 478 +1
Lines 71493 71543 +50
Branches 11825 11833 +8
==========================================
- Hits 62120 60318 -1802
- Misses 6239 8209 +1970
+ Partials 3134 3016 -118
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:
|
1343b06 to
3d8ce24
Compare
3d8ce24 to
b191894
Compare
167cf39 to
ee95f3c
Compare
|
Is this still in draft? Or is this ready for review? |
ee95f3c to
c2b77d2
Compare
|
/gemini review |
|
This PR is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you. |
tf.DType's __hash__/__eq__ are coerced to match plain Python ints (hash(tf.float32) == hash(1) == 1, tf.float32 == 1 is True). Once a tf.DType populated _DTYPE_CACHE or the result_type lru_cache, a later unrelated invalid int/bool argument would silently return the wrong cached dtype string instead of raising ValueError. Both caches are now keyed on (type(dtype), dtype) tuples, which defeats the coercing equality. Adds regression tests pinning the fix, extends the standardize_dtype contract test with a complex-dtype case, and adds a backend-agnostic test for the cache fast path.
The tf.DType coercion rationale was fully re-explained in three places across two files; keep one full explanation per file and shorten the others to what's not already said nearby. Drop a comment that only restated the line below it.
Moves the lru_cache-decorated resolution core above result_type, its sole call site, matching the house convention of defining private helpers before the public function that uses them.
Matches the house style used elsewhere in the codebase (e.g. backend/jax/rnn.py) instead of embedding an arrow mid-sentence.
Per review, this PR should cover `ALLOWED_DTYPES` and `standardize_dtype` only. The `result_type` lru_cache and its two tests are removed here and move to their own PR; `result_type` is restored to its master form. `ALLOWED_DTYPES_SET` stays because `standardize_dtype` is what needs the O(1) membership test.
variables.py had 14 comment lines on ~50 added lines, denser than the file's own standardize_shape (its closest sibling). Cut the "Fast path N" labels that just restated the branch below them, and folded the 6-line _DTYPE_CACHE block down to match the phrasing already settled on for the same explanation in dtypes.py. variables_test.py's new tests carried 6-line docstrings against a file where docstrings, when present, are one line. Trimmed to that convention and dropped inline comments that only restated the line below them.
f77140e to
c6593f1
Compare
keerthanakadiri
added
the
stat:awaiting keras-eng
| if isinstance(dtype, str): | ||
| if dtype in dtypes.ALLOWED_DTYPES_SET: | ||
| return dtype |
There was a problem hiding this comment.
How much speed-up to you get if you undo everything, but keep ALLOWED_DTYPES_SET, add these 3 lines of code...
hertschuh
added
stat:awaiting response from contributor
and removed
stat:awaiting keras-eng
Per review, undo the dtype cache and the restructured branches, keeping only the O(1) membership set and an early return for an already-canonical dtype string.
pctablet505
changed the title
google-ml-butler
Bot
added
kokoro:force-run
ready to pull
Removed comment about O(1) membership set for ALLOWED_DTYPES.
google-ml-butler
Bot
removed
the
ready to pull
Closes #23272
Summary
standardize_dtyperuns a chain ofhasattrprobes and a linear membership test against theALLOWED_DTYPEStuple on every call, even when it is handed a dtype string that is already canonical โ which is the overwhelmingly common case.Two changes:
dtypes.ALLOWED_DTYPES_SET, afrozensetderived fromALLOWED_DTYPES, so the membership test is O(1).ALLOWED_DTYPESitself stays an ordered tuple: several test suites build parametrized cases by iterating it, and pytest-xdist needs a deterministic collection order across workers.strthat is already in that set, skipping the normalization chain entirely.Everything else is unchanged.
Effect
Single process, interleaved with an A/A control, 20k iterations ร 11 rounds:
"float32"(canonical string)torch.float32(non-string path)3.4ร on the common path. The non-string path pays one extra
isinstancecheck and is ~3% slower; the A/A control puts the noise floor under 1%, so that cost is real but small.End to end,
standardize_dtypeis called 71 times per forward pass on a small CNN and 297 times on a 20-layer MLP (counted withcProfile, which is immune tofrom โฆ importrebinding). Measured PR vs master, CPU, medians of interleaved runs:The MLP figure is outside its noise band and agrees with the arithmetic (297 ร 77.7 ns = 23 ยตs on a 1.32 ms step). The CNN does proportionally more real compute per layer call, so the effect disappears into noise there.
Equivalence
A differential probe of 48 inputs โ every
ALLOWED_DTYPESstring, everyPYTHON_DTYPES_MAPkey,None, Python scalar types, numpy dtypes and scalar types, torch dtypes,"torch.*"/"jax.numpy.*"strings,"torch.float32.extra", and invalid inputs including1,True,[],{}andobject()โ returns identical results and identical exception types and messages on both trees. The only differing byte is the memory address inside theobject()error message, which differs between any two processes.Contributor Agreement