Repository navigation
Preserve promoted dtype in optimized batching - #1447
Open
DiadkoShmek wants to merge 1 commit into
Open
DiadkoShmek wants to merge 1 commit into
DiadkoShmek wants to merge 1 commit into
Conversation
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Shared-memory batching changes the dtype inferred by ordinary
Dataset.batch: an accepted batch containing int8[1]followed by int64[300]becomes int8[[1], [44]]afterenable_shared_memory_output(). The parallel-copy helper has the same first-element allocation assumption, including its non-shared output path.Infer the output dtype from every input leaf before allocating the output, matching
np.stackpromotion. Apply this at the three optimized allocation sites while retaining object/string fallback behavior.The regression uses real Map/Iter datasets and parallel-copy paths with both input orders, integer widths, bool/integer, integer/float and wider float/complex inputs. On the exact base 15 cases fail or error and 13 controls pass; all 28 pass after the fix. The selected supported batch suite passes 96 tests, including the new regressions. Separate property checks cover 288 ordered dtype/pipeline cases and noncontiguous inputs; independent real shared-memory/thread-copy checks also pass.
The unrestricted batch file retains the same 12 baseline errors: nine need optional JAX, and three are existing no-JAX tree-error message mismatches. No full-suite, multiprocess-prefetch, JAX/TF or throughput claim is made. Shared-memory tests close/unlink the allocations they own.
📚 Documentation preview 📚: https://google-grain--1447.org.readthedocs.build/