Allow more deeply nested mappings in params while retaining tight typing.

CLAUDE.md

@@ -276,6 +276,12 @@ initial_regimes = ["working", "working", "retired"]
 - Google-style docstrings
 - All functions require type annotations
 - Pre-commit hooks ensure code quality
+- Never use `from __future__ import annotations` — this project requires Python 3.14+
+
+### Testing Style
+
+- Use plain pytest functions, never test classes (`class TestFoo`)
+- Use `@pytest.mark.parametrize` for test variations
 
 ### Testing Style
 

src/lcm/params/__init__.py

@@ -0,0 +1,26 @@
+from collections.abc import Mapping, Sequence
+from typing import Any, overload
+
+from lcm.params.mapping_leaf import MappingLeaf
+from lcm.params.sequence_leaf import SequenceLeaf
+
+
+@overload
+def as_leaf(data: Mapping[str, Any]) -> MappingLeaf: ...
+
+
+@overload
+def as_leaf(data: Sequence[Any]) -> SequenceLeaf: ...
+
+
+def as_leaf(data: Mapping[str, Any] | Sequence[Any]) -> MappingLeaf | SequenceLeaf:
+    """Wrap a Mapping or Sequence as a JAX-pytree leaf."""
+    if isinstance(data, Mapping):
+        return MappingLeaf(dict(data))
+    if isinstance(data, Sequence):
+        return SequenceLeaf(data)
+    msg = f"as_leaf() expects a Mapping or Sequence, got {type(data).__name__}"
+    raise TypeError(msg)
+
+
+__all__ = ["MappingLeaf", "SequenceLeaf", "as_leaf"]

src/lcm/params/mapping_leaf.py

@@ -0,0 +1,46 @@
+"""A Mapping wrapper that is a JAX pytree but not itself a Mapping."""
+
+from collections.abc import Mapping
+from typing import Any
+
+import jax
+
+
+class MappingLeaf:
+    """A Mapping wrapper that is a JAX pytree but not itself a Mapping.
+
+    Prevents flatten_regime_namespace from recursing into contents while
+    allowing JAX to trace through array values.
+
+    Data is frozen to immutable containers on construction.
+    """
+
+    __slots__ = ("data",)
+
+    def __init__(self, data: Mapping[str, Any]) -> None:
+        from lcm.utils import ensure_containers_are_immutable  # noqa: PLC0415
+
+        self.data = ensure_containers_are_immutable(data)
+
+    def __repr__(self) -> str:
+        return f"MappingLeaf({dict(self.data)!r})"
+
+    __hash__ = None  # MappingProxyType is not hashable
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, MappingLeaf):
+            return NotImplemented
+        return self.data == other.data
+
+
+def _flatten(nmp: MappingLeaf) -> tuple[list[Any], tuple[str, ...]]:
+    keys = tuple(sorted(nmp.data.keys()))
+    values = [nmp.data[k] for k in keys]
+    return values, keys
+
+
+def _unflatten(keys: tuple[str, ...], values: list[Any]) -> MappingLeaf:
+    return MappingLeaf(dict(zip(keys, values, strict=True)))
+
+
+jax.tree_util.register_pytree_node(MappingLeaf, _flatten, _unflatten)

src/lcm/params/sequence_leaf.py

@@ -0,0 +1,45 @@
+"""A Sequence wrapper that is a JAX pytree but not itself a Sequence."""
+
+from collections.abc import Sequence
+from typing import Any
+
+import jax
+
+
+class SequenceLeaf:
+    """A Sequence wrapper that is a JAX pytree but not itself a Sequence.
+
+    Prevents flatten_regime_namespace from recursing into contents while
+    allowing JAX to trace through array values.
+
+    Data is frozen to immutable containers on construction.
+    """
+
+    __slots__ = ("data",)
+
+    def __init__(self, data: Sequence[Any]) -> None:
+        from lcm.utils import _make_immutable  # noqa: PLC0415
+
+        self.data = tuple(_make_immutable(v) for v in data)
+
+    def __repr__(self) -> str:
+        return f"SequenceLeaf({list(self.data)!r})"
+
+    def __hash__(self) -> int:
+        return hash(self.data)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SequenceLeaf):
+            return NotImplemented
+        return self.data == other.data
+
+
+def _flatten(sl: SequenceLeaf) -> tuple[list[Any], None]:
+    return list(sl.data), None
+
+
+def _unflatten(_aux: None, values: list[Any]) -> SequenceLeaf:
+    return SequenceLeaf(values)
+
+
+jax.tree_util.register_pytree_node(SequenceLeaf, _flatten, _unflatten)

src/lcm/typing.py

@@ -5,6 +5,9 @@
 from jax import Array
 from jaxtyping import Bool, Float, Int, Scalar
 
+from lcm.params import MappingLeaf
+from lcm.params.sequence_leaf import SequenceLeaf
+
 type ContinuousState = Float[Array, "..."]
 type ContinuousAction = Float[Array, "..."]
 type DiscreteState = Int[Array, "..."]
@@ -40,7 +43,17 @@
     bool
     | float
     | Array
-    | Mapping[str, bool | float | Array | Mapping[str, bool | float | Array]],
+    | MappingLeaf
+    | SequenceLeaf
+    | Mapping[
+        str,
+        bool
+        | float
+        | Array
+        | MappingLeaf
+        | SequenceLeaf
+        | Mapping[str, bool | float | Array | MappingLeaf | SequenceLeaf],
+    ],
 ]
 
 # Internal regime parameters: A flat mapping with function-qualified names.

src/lcm/utils.py

@@ -9,13 +9,17 @@
 from dags.tree import flatten_to_qnames, unflatten_from_qnames
 from jax import Array
 
+from lcm.params import MappingLeaf
+from lcm.params.sequence_leaf import SequenceLeaf
 from lcm.typing import RegimeName
 
 T = TypeVar("T")
 
 
 def _make_immutable(value: Any) -> Any:  # noqa: ANN401
     """Recursively convert a value to its immutable equivalent."""
+    if isinstance(value, (MappingLeaf, SequenceLeaf)):
+        return value  # already immutable by construction
     if isinstance(value, (MappingProxyType, tuple, frozenset)):
         return value
     if isinstance(value, Mapping):
@@ -51,8 +55,12 @@ def ensure_containers_are_immutable[K, V](
     return cast("MappingProxyType[K, V]", _make_immutable(value))
 
 
-def _make_mutable(value: Any) -> Any:  # noqa: ANN401
+def _make_mutable(value: Any) -> Any:  # noqa: ANN401, PLR0911
     """Recursively convert a value to its mutable equivalent."""
+    if isinstance(value, MappingLeaf):
+        return {k: _make_mutable(v) for k, v in value.data.items()}
+    if isinstance(value, SequenceLeaf):
+        return [_make_mutable(v) for v in value.data]
     if isinstance(value, (set, list)):
         return value
     if isinstance(value, (MappingProxyType, Mapping)):