Drop _lcm/jaxtyping_patch.py in favour of upstream fix#365
Open
hmgaudecker wants to merge 1 commit into
Open
Conversation
jaxtyping's `_array_types` module-level `object()` sentinels broke under cloudpickle round-trips because plain `object()` does not survive serialisation with identity intact. The workaround lived as a local monkey-patch in `_lcm/jaxtyping_patch.py` for the variadic-dim sentinel. Replace the workaround by depending on a jaxtyping branch that fixes all three affected sentinels (`_any_dtype`, `_anonymous_dim`, `_anonymous_variadic_dim`) at the source via `__reduce__`-backed singleton classes (patrick-kidger/jaxtyping#390 — pending upstream review). Floor pin bumps to `jaxtyping>=0.3.10`; the pixi pypi-dependency override pulls the fork branch until upstream releases. `tests/simulation/test_simulate_aot.py` covers the cloudpickle round-trip paths and continues to pass without the local patch.
Documentation build overview
34 files changed ·
|
Benchmark comparison (main → HEAD)Comparing
|
5 tasks
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.
Summary
Removes the local monkey-patch at
src/_lcm/jaxtyping_patch.pyin favour of an upstream fix in jaxtyping that covers all three affected sentinels (_any_dtype,_anonymous_dim,_anonymous_variadic_dim), not just the variadic-dim sentinel we patched locally.Until the upstream fix is released as
jaxtyping >= 0.3.10, pylcm pulls the patched branch directly via a pixipypi-dependenciesgit override — mirroring the pattern in place fordags. Once the release lands, the override and the temporary comment can both go.Changes
pyproject.toml:jaxtyping>=0.3.2 → jaxtyping>=0.3.10; pixipypi-dependenciesadds the git override pointing athmgaudecker/jaxtyping@fix/sentinel-cloudpickle.src/_lcm/jaxtyping_patch.py.src/_lcm/__init__.py— dropped the load-bearing-import comment block, sincejaxtyping_patchwas the only thing needing pre-emption.pixi.lockfor all four pixi envs.Background
jaxtyping's
_array_typesmodule-levelobject()sentinels broke under cloudpickle round-trips because plainobject()does not survive serialisation with identity intact. The variadic-dim case is what beartype-claw'd pylcm hit through DAG-built cloudpickleable callables; the other two sentinels would fail similarly under cloudpickle, just through different paths. The upstream fix replaces all three with__reduce__-backed singleton classes.Upstream issue: patrick-kidger/jaxtyping#390. Fix branch:
hmgaudecker/jaxtyping@fix/sentinel-cloudpickle(upstream PR pending).Test plan
pixi run -e type-checking ty— clean.prek run --all-files— clean.pixi run --environment tests-cpu tests -n 4— 1018 passed, 42 skipped. The cloudpickle round-trip tests intests/simulation/test_simulate_aot.pyare the binding ones for what the local patch was protecting.Stacking
Stacked on #363. Rebase / change base to
mainonce #363 lands; the diff is independent of the mahler-yum example.