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

src/quax/_compat.py

@@ -134,9 +134,29 @@ def scan_bind_params(
         num_xs: int,
         num_ys: int,
     ) -> dict[str, Any]:
-        """Parameters for re-binding `scan_p` with the given group sizes."""
-        del num_xs, num_ys
-        return {**params, "num_consts": num_consts, "num_carry": num_carry}
+        """Parameters for re-binding `scan_p` with the given group sizes.
+
+        `linear` is a per-operand flag tuple (consts + carry + xs), dropped from
+        `scan_p` in JAX 0.10.2. Quaxifying the body can change the flat operand
+        count -- a single `ArrayValue` may flatten to several arrays -- in which
+        case the incoming `linear` no longer lines up with the new operands and
+        JAX's scan rules raise.
+
+        Only rebuild `linear` when it is present *and* the operand count actually
+        changed. When it matches (the common single-leaf case), keep JAX's
+        original linearity analysis so `lax.scan`'s AD does not lose it to an
+        all-`False` rebuild; when JAX no longer takes the parameter, never
+        reintroduce it.
+        """
+        del num_ys
+        new_params = {**params, "num_consts": num_consts, "num_carry": num_carry}
+        n_operands = num_consts + num_carry + num_xs
+        linear = params.get("linear")
+        if linear is not None and len(linear) != n_operands:
+            # Operand count changed: the old per-operand flags no longer align.
+            # The body is retraced from scratch, so no linearity carries over.
+            new_params["linear"] = (False,) * n_operands
+        return new_params
 
 
 typeof: Callable[[Any], Any]

src/quax/examples/unitful/_core.py

@@ -16,6 +16,15 @@ def __init__(self, name):
     def __repr__(self):
         return self.name
 
+    def __lt__(self, other):
+        # `units` is a dict keyed by `Dimension`, and JAX sorts dict keys when
+        # flattening a pytree -- which equinox does to every `Module.__init__`
+        # argument. Without an ordering, any multi-dimension `Unitful` (e.g.
+        # `{meters: 1, seconds: -1}`) fails to construct.
+        if not isinstance(other, Dimension):
+            return NotImplemented
+        return self.name < other.name
+
 
 kilograms = Dimension("kg")
 meters = Dimension("m")