Skip to content

Clarify fake_pmap_and_jit flags and warn on positional use - #478

Open
sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/144-clear-fake-pmap-jit-flags-20261011
Open

sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/144-clear-fake-pmap-jit-flags-20261011

Conversation

@sylvesterkaczmarek

Copy link
Copy Markdown

Fixes #144.

Changes

  • Add explicit keyword-only disable_pmap and disable_jit flags to chex.fake_pmap_and_jit so callers can see exactly which JAX transformations are disabled.
  • Keep the historical enable_pmap_patching and enable_jit_patching keywords and the two legacy positional flags backward compatible.
  • Emit a DeprecationWarning for positional flags, particularly the dangerous single-flag form that silently leaves JIT disabled with its default true value.
  • Reject conflicting new/old flag names and duplicate positional/keyword arguments with actionable TypeErrors.
  • Expand docstrings with an example showing how to set both flags from a debug setting.
  • Add parameterized regressions for all four independent patching combinations, legacy warning behavior and conflicting arguments.

Validation

  • Fake-context and variants test suites: 144 passed, 25 skipped.
  • Ruff and Python syntax checks passed; git diff --check passed.
  • Public legacy keyword behavior remains covered by the existing test matrix.

No changes to jax.jit or jax.pmap themselves. This makes the old positional API noisy but retains its behavior to avoid breaking existing call sites.

@sylvesterkaczmarek
sylvesterkaczmarek force-pushed the fix/144-clear-fake-pmap-jit-flags-20261011 branch from bcb6c32 to debd4d0 Compare October 11, 2026 12:33
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.

fake_pmap_and_jit has a confusing interface

1 participant