[consts] Fix handling of closed-over constants in MPMD - #1404
Open
copybara-service[bot] wants to merge 1 commit into
Open
[consts] Fix handling of closed-over constants in MPMD#1404copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into
Conversation
copybara-service
Bot
force-pushed
the
test_929077635
branch
from
June 9, 2026 11:02
5471d2f to
90e9a04
Compare
These changes are needed when we will set JAX_USE_SIMPLIFIED_JAXPR_CONSTANTS=1, to enable the new JAX handling of closed-over constants. Until then, `num_const_args` is 0 throughout. See https://docs.jax.dev/en/latest/internals/constants.html for some implementation details. One challenge with handling closed-over constants is that jax_fn_info (a sequence of ArgInfo) refers strictly to the actual function arguments, while the argument lists include the hoisted constant arguments as a prefix (`num_const_args` of them). In this CL we adjust the code to handle the case when `num_const_args` is not 0. PiperOrigin-RevId: 929077635
copybara-service
Bot
force-pushed
the
test_929077635
branch
from
June 9, 2026 11:26
90e9a04 to
27a4c25
Compare
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
1 participant
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.
[consts] Fix handling of closed-over constants in MPMD
These changes are needed when we will set JAX_USE_SIMPLIFIED_JAXPR_CONSTANTS=1,
to enable the new JAX handling of closed-over constants. Until then,
num_const_argsis 0 throughout.See https://docs.jax.dev/en/latest/internals/constants.html for some implementation details.
One challenge with handling closed-over constants is that jax_fn_info
(a sequence of ArgInfo) refers strictly to the actual function arguments,
while the argument lists include the hoisted constant arguments as
a prefix (
num_const_argsof them).In this CL we adjust the code to handle the case when
num_const_argsis not 0.