Backport PR #205 on branch versions/v0.3.x (🐛 fix(compat): don't reintroduce scan_p's removed linear param)

[nstarman]

Manual backport of #205 (MrMeeseeks reported a conflict).

Conflict resolution

versions/v0.3.x never had the linear handling at all — the rebuild was introduced on main by #189 ("support multi-leaf Values and empty scans"), which was not backported. So the conflict is main's two-commit history vs. this branch's plain pass-through.

Resolved by taking the net effect of #205 on scan_bind_params only:

• rebuild linear when it is present and the operand count changed (the fix: scan support for multi-leaf Values and length-0 scans #189 fix this branch lacks — quaxifying can change the flat operand count, and a stale per-operand tuple makes JAX's scan rules raise);
• never inject linear when JAX no longer takes it (the 🐛 fix(compat): don't reintroduce scan_p's removed linear param #205 fix; JAX 0.10.2 dropped the parameter).

The cherry-pick would also have dragged in unrelated main-only changes to _compat.py (import reordering, dropping the JAX_GE_0_7_0 / pjit_p fallback, comment rewraps). Those are excluded — this branch still supports jax>=0.5.3 and Python 3.10, where the pjit_p fallback is required.

Note on scope

This branch was not hit by the weekly-CI failure that motivated #205: its pass-through never added linear, so TypeError: _scan_abstract_eval() got an unexpected keyword argument 'linear' does not occur here (verified — scan tests pass on versions/v0.3.x @ b365e08 with jax==0.10.2). The backport keeps the two branches consistent and brings the operand-count-misalignment guard across.

Verification

All on Python 3.11:

• jax==0.10.2, scan tests: 9 passed.
• jax==0.10.2, full suite: 1 failed, 967 passed, 137 skipped, 58 xfailed, 6 xpassed. The one failure, tests/usage/test_unitful.py::test_integer_pow_units_and_values_agree, is pre-existing on b365e08 with this JAX (confirmed on the unmodified base) and unrelated to scan.
• jax==0.6.2 (exercises the linear-present path), scan tests: 9 passed.

🤖 Generated with Claude Code

[nstarman]

Fixed the CI failure in 3d011d0.

It was the pre-existing tests/usage/test_unitful.py::test_integer_pow_units_and_values_agree failure noted above, failing on every job (all Python versions, oldest and newest deps) — not the scan change:

ValueError: Comparator raised exception while sorting pytree dictionary keys.
  └── TypeError: '<' not supported between instances of 'Dimension' and 'Dimension'

Equinox flattens every Module.__init__ argument, and JAX sorts pytree dict keys while flattening. Unitful.units is a dict[Dimension, int], so any multi-dimension value ({meters: 1, seconds: -1}) can't be constructed. Single-dimension values are fine — a 1-element sort makes no comparisons — which is why only this one test trips it. main doesn't hit it because its _FastModuleMeta.__call__ bypasses that equinox flatten.

Fix: Dimension.__lt__ comparing names, so the sort has a total order.

Verified on Python 3.11 / jax==0.10.2: full suite 968 passed, 137 skipped, 58 xfailed, 6 xpassed (was 1 failed / 967 passed). Also passes on --resolution lowest-direct / Python 3.10.