[nstarman]

JAX 0.11.2 registers the built-in frozendict (PEP 814) as a PyTree and flattens it like dict: leaves in sorted-key order, unflatten rebuilds in sorted order, and keys that can't be ordered raise an error. ImmutableMap flattened in insertion order, so switching to the built-in on Python 3.15+ silently reordered leaves. This PR makes ImmutableMap.tree_flatten sort its keys so the two types agree.

⚠️ Breaking

• jax.tree.leaves(ImmutableMap(b=1, a=2)) was [1, 2] and is now [2, 1]. This affects code that indexes leaves, saved optimizer states, and checkpoints.
• After a tree.map / jit round-trip, iteration order is now sorted rather than insertion order. Equality and hashing are already order-independent.
• Maps with keys that can't be ordered against each other (e.g. {1: …, "b": …}) now raise TypeError on flatten. Native JAX raises ValueError for the same case on dict and frozendict. I didn't wrap it to match.

The treedefs still differ (CustomNode(ImmutableMap[...]) vs frozendict({...})), and keypaths are unchanged.

Changes

• src/xmmutablemap/_core.py: tree_flatten sorts its keys. The docstring and doctest are updated.
• tests/test_immutablemap.py: new tests for sorted leaves, the sorted round-trip, and unorderable keys.
• README.md: the leaf-order warning added in 📝 docs: prefer the built-in frozendict on Python 3.15+ #79 is rewritten to say ImmutableMap matches the built-in and notes the change from earlier releases.

Verification

• pytest passes on 3.10 and 3.13 (116 tests). prek run --all-files passes, including mypy.
• On CPython 3.15.0 with jax 0.11.2, ImmutableMap(b=1, a=2, c=3) and frozendict(b=1, a=2, c=3) give identical jax.tree.leaves ([2, 1, 3]) and identical round-tripped key order.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (main@091e220). Learn more about missing BASE report.

Additional details and impacted files

@@           Coverage Diff            @@
##             main       #80   +/-   ##
========================================
  Coverage        ?   100.00%           
========================================
  Files           ?         3           
  Lines           ?        75           
  Branches        ?         0           
========================================
  Hits            ?        75           
  Misses          ?         0           
  Partials        ?         0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.