xmmutablemap

JAX-compatible Immutable Mapping

JAX prefers immutable objects, but before Python 3.15 there is no built-in immutable dictionary. 😢
This repository defines a light-weight immutable map (lower-level than a dict) that JAX understands as a PyTree. 🎉 🕶️ On Python 3.15+, prefer the built-in frozendict (see below).

Installation

pip install xmmutablemap

uv add xmmutablemap

pip install git+https://github.com/GalacticDynamics/xmmutablemap.git

cd /path/to/parent
git clone https://github.com/GalacticDynamics/xmmutablemap.git
cd xmmutablemap
pip install -e .  # editable mode

Documentation

xmutablemap provides the class ImmutableMap, which is a full implementation of Python's Mapping ABC. If you've used a dict then you already know how to use ImmutableMap! The things ImmutableMap adds is 1) immutability (and related benefits like hashability) and 2) compatibility with JAX.

from xmmutablemap import ImmutableMap

print(ImmutableMap(a=1, b=2, c=3))
# ImmutableMap({'a': 1, 'b': 2, 'c': 3})

print(ImmutableMap({"a": 1, "b": 2.0, "c": "3"}))
# ImmutableMap({'a': 1, 'b': 2.0, 'c': '3'})

Relationship to frozendict (PEP 814)

Python 3.15 adds a built-in frozendict (PEP 814), which covers most of what ImmutableMap was written for. ImmutableMap is now a subclass of it, so the two interoperate directly:

from xmmutablemap import ImmutableMap, frozendict

d = ImmutableMap(a=1, b=2)

print(isinstance(d, frozendict))
# True
print(d == frozendict(a=1, b=2) == {"a": 1, "b": 2})
# True

On Python 3.15+, xmmutablemap.frozendict is the built-in. On older Pythons it is a behaviour-matched stand-in, so this code works the same on every supported version and you can migrate ahead of your Python floor.

Prefer the built-in frozendict on Python 3.15+

Since JAX 0.11.2, JAX registers the built-in frozendict as a PyTree node. On Python 3.15+ with jax>=0.11.2, use the built-in frozendict directly; it does everything ImmutableMap does, with no extra dependency. ImmutableMap remains for Python < 3.15, where the built-in does not exist (and JAX treats xmmutablemap.frozendict, the stand-in, as an opaque leaf).

import jax

print(jax.tree.structure(ImmutableMap(a=1, b=2)))
# PyTreeDef(CustomNode(ImmutableMap[('a', 'b')], [*, *]))
print(jax.tree.structure(frozendict(a=1, b=2)))
# Python 3.15+, jax>=0.11.2: PyTreeDef(frozendict({'a': *, 'b': *}))
# older Pythons (the stand-in): PyTreeDef(*)

One deliberate difference: ImmutableMap | other merges any Mapping, while frozendict | other merges only a dict or frozendict (so, for example, a custom Mapping with no | of its own raises TypeError). Convert such mappings with dict(...) when switching to the built-in.

Subclassing

Being a frozendict, an ImmutableMap gets its contents in __new__, and they cannot change afterwards. A subclass that wants to transform its input overrides __new__; its __init__, if any, runs on an already-filled map. Subclasses are not PyTree nodes unless they register themselves:

from jax.tree_util import register_pytree_node_class


@register_pytree_node_class
class UpperKeys(ImmutableMap):
    def __new__(cls, *args, **kwargs):
        data = dict(*args, **kwargs)
        return super().__new__(cls, {k.upper(): v for k, v in data.items()})


print(UpperKeys(a=1))
# UpperKeys({'A': 1})
print(jax.tree.leaves(UpperKeys(a=1, b=2)))
# [1, 2]

JAX Integration

One of the key benefits of ImmutableMap is its compatibility with JAX. Since it's immutable and hashable, it can be used in places where JAX would normally complain about mutable objects like regular dictionaries.

Using ImmutableMap as a Default in JAX Dataclasses

Here's an example showing how ImmutableMap can be used as a default value in a dataclass, which is particularly useful with JAX:

import functools
import jax
import jax.numpy as jnp
from dataclasses import dataclass
from xmmutablemap import ImmutableMap


@functools.partial(
    jax.tree_util.register_dataclass, data_fields=["params"], meta_fields=["batch_size"]
)
@dataclass(frozen=True)
class Config:
    """Configuration with immutable default parameters."""

    # This works! ImmutableMap is immutable and hashable
    params: ImmutableMap[str, float] = ImmutableMap(
        learning_rate=0.001, momentum=0.9, weight_decay=1e-4
    )
    batch_size: int = 32


# JAX can safely transform functions using this dataclass
@jax.jit
def train_step(config: Config, data: jnp.ndarray) -> jnp.ndarray:
    """Example training step that uses config parameters."""
    lr = config.params["learning_rate"]
    return data * lr


# This works perfectly
config = Config()
data = jnp.array([1.0, 2.0, 3.0])
result = train_step(config, data)
print(f"Result: {result}")
# Result: [0.001 0.002 0.003]

Key Benefits for JAX

• Immutability: Once created, ImmutableMap cannot be modified, preventing accidental mutations that could break JAX's functional programming model
• Hashability: JAX can safely cache and memoize functions that use ImmutableMap instances
• PyTree Support: ImmutableMap is registered as a JAX PyTree, so it works seamlessly with JAX transformations like jit, grad, vmap, etc.
• Safe Defaults: Can be used as default values in dataclasses without the typical pitfalls of mutable defaults

Development

We welcome contributions!