feat: register_pytree_node — allow custom classes in mx.compile

[st-adam]

Fixes #3499.

Addresses review feedback from @zcbenz: implementation moved to C++ so mx.compile is natively aware of registered pytree types. The previous Python-side utils.compile wrapper has been removed.

Summary

• New mx.register_pytree_node(cls, flatten_fn, unflatten_fn) API (mirrors jax.tree_util.register_pytree_node)
• C++ registry in python/src/trees.cpp, exposed via init_trees() nanobind module
• PyCompiledFun::call_impl recurse handles registered types: their type-id + aux hash participate in the compile cache key
• Existing tree traversal (tree_visit, tree_map, tree_visit_update) recurse into registered nodes — same code path used by tree_unflatten on the compile path
• mlx.utils re-exports register_pytree_node for the natural from mlx.utils import … access pattern

API

import mlx.core as mx

class Pair:
    def __init__(self, a, b):
        self.a = a
        self.b = b

mx.register_pytree_node(
    Pair,
    lambda p: ([p.a, p.b], None),          # flatten: (children, aux_data)
    lambda _aux, children: Pair(*children), # unflatten
)

@mx.compile
def add_pair(p):
    return p.a + p.b

add_pair(Pair(mx.array(3), mx.array(4)))    # → array(7, dtype=int32)

Test plan

python/tests/test_compile.py::TestCompile::test_compile_registered_pytree_node covers:

• Pre-registration rejection (unchanged error path)
• Compiled forward with a registered custom class
• aux_data participation in cache key (tagged subclass → distinct trace)
• Malformed flatten_fn return surfaces a clean ValueError

Regression sweep on the CPU build:

• python/tests/test_compile.py — 55 passed
• python/tests/test_tree.py — 4 passed
• python/tests/test_autograd.py + test_vmap.py — full suites green

Implementation notes

• The C++ registry is heap-allocated and never freed. A function-local std::unordered_map<PyTypeObject*, PytreeNodeDef> triggers a use-after-finalize segfault at interpreter shutdown because the stored nb::callables outlive Python state. This matches the lifetime trick used by structure_sentinel().
• The compile-cache fingerprint mixes id(type) with hash(aux_data) (golden-ratio mixing constant) so two structurally distinct registered instances retrace correctly. Unhashable aux falls back to type-only fingerprinting.
• flatten_fn may return children as either list or tuple; the registry rejects non-sequence return values up front.

Files

• python/src/trees.h (+38) — public API + helpers
• python/src/trees.cpp (+263) — registry, init_trees, tree-traversal hooks
• python/src/transforms.cpp (+15) — PyCompiledFun::call_impl recurse + error-message hint
• python/src/mlx.cpp (+2) — wire init_trees() into NB_MODULE
• python/mlx/utils.py (+1 net) — re-export register_pytree_node
• python/tests/test_compile.py (+63) — coverage

🤖 Generated with Claude Code

[zcbenz]

I think the register_pytree_node API is a good thing to have, but utils.compile would be a bad addition.

The register_pytree_node API should be implemented in C++ layer so mx.compile can be made aware of it directly. And we should probably remove the python versions of tree utils and expose the C++ ones instead so we don't have to duplicate the register_pytree_node implementation in 2 languages.

[st-adam]

@zcbenz Thanks for the review. I've moved the implementation into C++ as requested:

• Registry, register_pytree_node, and PyCompiledFun::call_impl recurse all live in python/src/trees.{cpp,h} and python/src/transforms.cpp.
• The Python utils.compile wrapper is gone; mx.compile natively handles registered types now.
• mlx.utils.register_pytree_node is a thin re-export of the C++ binding, so we keep a single source of truth as you suggested.
• Existing tree_visit / tree_map / tree_visit_update overloads recurse into registered nodes, which is what tree_unflatten rides on the compile path — so the trace/unflatten round-trip just works.
• I left the Python tree_* utilities in python/mlx/utils.py untouched for this PR. Migrating them to C++ bindings feels like a separate, larger change; happy to do that as a follow-up if you'd prefer it bundled.

Verified locally on a CPU build: test_compile.py (55), test_tree.py (4), test_autograd.py + test_vmap.py all green, plus a new test_compile_registered_pytree_node covering the registered-pair, aux-data retrace, and malformed-flatten cases.

[zcbenz]

This declaration is not needed, already declared in python/src/mlx.cpp.

[zcbenz]

You can use nb::type_object and let nanobind do the check.

[zcbenz]

We should probably just use PyObject* as registry key, force converting to PyTypeObject* adds code complexity without much benefits.

[zcbenz]

flatten_registered and unflatten_registered are helpers that do not need to be exposed in header.

[zcbenz]

You can use nb::hash

[zcbenz]

Can you just do auto children = nb::cast<std::vector<nb::object>>(seq[0]);?

[zcbenz]

The nb::cast<nb::sequence> should be able to do type check so I think this check is redundant.

[zcbenz]

There is no need to be defensive here since it is inside a try/catch, just do castings and let it throw when bad happens.

[zcbenz]

You can just pass subtree here?

[st-adam]

@zcbenz Thanks for the detailed review — all nine inline comments addressed in 2f7c5e6:

• trees.h surface is now register_pytree_node, is_registered_pytree, pytree_children, registered_pytree_fingerprint. flatten_registered / unflatten_registered moved into the anonymous namespace in trees.cpp. Stale init_trees declaration removed.
• register_pytree_node takes nb::type_object cls so nanobind enforces the type check.
• Registry is now keyed by PyObject* directly — no PyTypeObject* reinterpret cast.
• flatten_registered uses nb::cast<std::vector<nb::object>>(seq[0]) and lets nb::cast<nb::sequence> enforce the (children, aux) shape.
• Fingerprint uses nb::hash and lets nanobind throw on unhashable aux.
• tree_visit_update / tree_map pass the subtree handle directly to unflatten_registered (no fabricated type handle).
• transforms.cpp now uses the new pytree_children helper.

Lint also fixed (clang-format + black) in the previous commit. Local CPU run: test_compile.py (56), test_tree.py (4), test_autograd.py + test_vmap.py (64) all green. Net diff: -49 lines.

[angeloskath]

Sorry for being late to this but why do we think that register_pytree_node is a good addition? I find that specifically this option is often abused making modules be typed PyTrees and optimizers be typed PyTrees for very little added benefit. Do we have any real world use case where this simplifies the code? A good example that I would personally avoid is Equinox. Perhaps the problem there is not the proliferation of typed PyTrees but rather their immutability but I am not sure I see how

@mx.compile
def add_pair(p):
    return p.a + p.b

is better than

@mx.compile
def _add_pair(a, b):
    return a + b

def add_pair(p):
    return _add_pair(p.a, p.b)

A more real world example of this is simply the training loop. We don't need the model to be a typed PyTree, we can very quickly make the whole call "pure functional" by passing in a dictionary of parameters

@mx.compile
def step(params, opt_state, batch):
    model.update(params)
    loss_val, grads = nn.value_and_grad(model, loss_fn)(batch)
    optimizer.state = opt_state
    optimizer.update(model, grads)
    return (loss_val, model.parameters(), optimizer.state)

Changing gears after discussing whether we should add this at all, I see two issues to be addressed in the code

1. The pytree check should be first otherwise any custom type that subclasses tuple, list or dict cannot be a registered pytree.
2. In tree_visit_update we actually edit in-place the passed in tree. The registered types are now immutable like tuples which may be not the wanted behavior. For instance if I were to use a Pair from your example above in @partial(mx.compile, outputs=pair) it wouldn't quite work.

[st-adam]

Thanks @angeloskath. Over-complexity concern noted, and I'm happy to make this as simple as the use case actually requires. Before discussing implementation shape, let me lay out the problem this PR is trying to solve, since I think that context is what makes the API question concrete.

mx.compile only accepts function arguments that are trees of mx.arrays, scalars, lists, tuples, or dicts. Any other Python class is rejected at trace time with:

ValueError: [compile] Function arguments must be trees of arrays or constants ...
received type mlx_lm.models.cache.ArraysCache

Every model in mlx-lm that uses a cache class other than the plain KVCache hits this. Layer forward passes look like layer(x, cache=cache) where cache is one of ArraysCache, MambaCache, DeepseekV4Cache, etc. Those classes wrap mx.arrays inside a Python object so the layer can mutate state across decode steps. Compile sees the wrapper, can't trace through it, and rejects the call.

17 mlx-lm model architectures use these wrapped caches and therefore cannot be compiled today: qwen3_5, qwen3_next, lfm2, lfm2_moe, mamba, mamba2, jamba, plamo2, recurrent_gemma, rwkv7, nemotron_h, kimi_linear, longcat_flash_ngram, bailing_moe_linear, baichuan_m1, granitemoehybrid, falcon_h1. Plus the mlx-vlm wrappers that reuse them (qwen3_5, lfm2_vl) and any downstream serving stack that turns on JIT. In vMLX the symptom is that _apply_jit_compilation() has an explicit ArraysCache skip-block: the warmup forward raises, the engine catches it, logs JIT: warmup raised … rolling back mx.compile, and the model serves with JIT off for the rest of the process lifetime. Users opted into --enable-jit and silently get no fusion benefit on any hybrid SSM model. The same applies to anyone using @mx.compile on these models directly outside vMLX.

On workarounds. The pure-functional pattern (step(params, opt_state, batch)) is the right answer for Module / Optimizer, and I agree on the Equinox concern there. It doesn't compose with these caches because the cache mutates inside layer.__call__. To make compile happy without an opaque-wrapper escape hatch, the unwrap has to happen inside the compiled function, which means every one of those 17 models would have to be rewritten to take raw (keys, values, lengths) array lists per layer instead of cache objects. The wrapper pattern (def fwd(c): return _fwd(c.a, c.b)) is fine at user-code level but doesn't help here because the user doesn't see the cache; it lives behind the layer abstraction.

The current implementation in this PR is a C++ registry plus a public mx.register_pytree_node(cls, flatten_fn, unflatten_fn) API, mirroring jax.tree_util.register_pytree_node. Tree traversal and PyCompiledFun::call_impl recurse into registered types. Two technical issues you flagged still need fixes here: registered check must run before the isinstance(tree, (list, tuple, dict)) branches, and tree_visit_update should rebind via unflatten_registered so registered types behave as immutable.

The simpler alternative is to drop mx.register_pytree_node entirely (no registry, no public API) and have the class opt in by implementing two methods, same as __iter__ / __len__:

class ArraysCache:
    def tree_flatten(self) -> tuple[list, Any]:
        return self.cache, (self.size,)
    @classmethod
    def tree_unflatten(cls, aux, children):
        c = cls(*aux); c.cache = children; return c

MLX side: tree traversal + PyCompiledFun::call_impl recurse check hasattr(obj, "tree_flatten") and hasattr(type(obj), "tree_unflatten") before the isinstance(tree, (list, tuple, dict)) branch. Call the methods, recurse on children, fingerprint on aux. No global registry, no register_pytree_node API surface to abuse, and trees.cpp shrinks from ~263 to ~60 lines. The check-ordering issue is resolved by construction (protocol must run first to catch tuple/list/dict subclasses). The tree_visit_update mutation issue is resolved by construction too: tree_unflatten returns a new object, parent rebinds, matches immutable-tuple semantics, and @partial(mx.compile, outputs=pair) works because compile rebinds the parent. The trade-off is that you can't opt in classes you don't own; workaround is a subclass that implements the methods. For mlx-lm caches (the motivating case) we own the classes, so this isn't a constraint in practice.

Either shape works for me. I can pivot to the simpler alternative and drop the registry, or stay on the current implementation and land the two fixes you flagged. Whichever you and @zcbenz prefer, let me know and I'll push the change.