Skip to content

Reject unsupported named static arguments for pmap test variants - #477

Open
sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/115-reject-pmap-static-argnames-20261010
Open

sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/115-reject-pmap-static-argnames-20261010

Conversation

@sylvesterkaczmarek

Copy link
Copy Markdown

Fixes #115.

Summary

  • Fail clearly when static_argnames is supplied to the Chex with_pmap variant. JAX pmap cannot mark keyword arguments static and the option was previously silently ignored.
  • Explain how to use positional static_argnums or disable the pmap variant.
  • Preserve the behavior of empty or omitted static_argnames, as well as the supported with_jit variant.
  • Add three regression cases for tuple, string and empty named-static arguments.

Validation

  • python -m pytest -q chex/_src/variants_test.py: 102 passed, 25 skipped (device-dependent pmap tests).
  • ruff check chex/_src/variants.py chex/_src/variants_test.py: passed.
  • git diff --check: passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

chex.variants(with_pmap=True) ignores static_argnames

1 participant