Fix (critical!) variable ordering bug

brentyi · Sep 13, 2024 · bb32b1c · bb32b1c
1 parent 1927861
commit bb32b1c
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/jaxls/_variables.py b/src/jaxls/_variables.py
@@ -47,7 +47,9 @@ def ordered_dict_items[T](
         self,
         var_type_mapping: dict[type[Var[Any]], T],
     ) -> list[tuple[type[Var[Any]], T]]:
-        return sorted(var_type_mapping.items(), key=lambda x: x[0])
+        return sorted(
+            var_type_mapping.items(), key=lambda x: self.order_from_type[x[0]]
+        )
 
 
 @jdc.pytree_dataclass