Repository navigation
Preserve residual keyword arguments in fake_pmap - #484
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 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.
chex.fake_pmapcan drop keyword-only values,**kwargsand arguments after an omitted default. It binds the call but forwards onlyBoundArguments.args, so valid native-pmap calls either raise or silently use a default. In a two-replica SGD check, native pmap updates[2, -1]to[1.25, -1.375], while the base fake path silently leaves the parameters unchanged.Forward both positional and residual keyword arguments. Retain kwargs through the static-argument wrapper and add the synthetic parallel axis to both argument trees. Public transform signatures are unchanged.
All 12 new regression cases fail on exact base and pass with the fix. The selected fake/variants suite passes 170 tests, including those regressions, on two CPU devices. A separate closed-form least-squares gradient/pmean check agrees with native pmap and fixed fake execution in eager/JIT modes. Independent checks also cover default gaps, static Python branches,
None, trees/collectives, explicit axes, synthetic axes, native-invalid cases and context restoration.Validation uses real CPU JAX package imports; no pmap/math stand-ins or GPU/distributed-performance claim. Nearby open out_axes/static-argument/input-axis proposals remain separate from this dropped-kwargs fix.