feat(utils): add register_pytree_node and pytree-aware compile wrapper

[st-adam]

Fixes #3499.

Summary

• Adds mlx.utils.register_pytree_node(cls, flatten_fn, unflatten_fn) — a JAX-style API for registering custom classes as pytree nodes
• Extends tree_map and tree_flatten to traverse registered types
• Adds mlx.utils.compile() — a Python-level pytree-aware wrapper around mlx.core.compile that flattens registered-type arguments to arrays before the C++ tracer sees them

Motivation

mx.compile rejects custom objects as function arguments even when they are simple array wrappers. This blocks JIT compilation of hybrid SSM+attention models (Qwen 3.5/3.6, Llama 4, Gemma 3n, etc.) whose forward pass takes ArraysCache (SSM state cache) arguments alongside regular KVCache objects.

Usage

import mlx.utils as mlxu

# Register ArraysCache as a pytree
mlxu.register_pytree_node(
    ArraysCache,
    lambda c: (list(c.state), None),         # flatten → list of mx.array
    lambda _, s: ArraysCache.from_state(s),  # unflatten
)

# Use mlx.utils.compile (not mlx.core.compile) to get pytree support
compiled_forward = mlxu.compile(model.__call__)
output = compiled_forward(x, cache=cache)   # cache may contain ArraysCache

Test plan

• register_pytree_node + tree_flatten on custom class: returns correct flat array list
• mlx.utils.compile with registered pytree arg: correct result (3 + 4 + 10 = 17 verified)
• No-registry fast path: delegates directly to mlx.core.compile, zero overhead
• Structure-change detection: triggers recompile on next call

Limitations

SSM caches with in-place state mutation require the cache to also flow through inputs/outputs for the mutations to be captured by the tracer. This PR addresses argument-position structured inputs; mutable SSM state is a follow-on.

Files changed

• python/mlx/utils.py — all changes (pure Python, no C++ modifications)

🤖 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.