Support namedtuple and tuple subclasses in tree_merge - #3703

Summary

tree_map and tree_map_with_path reconstruct list/tuple containers with

return TreeType(*subtrees) if hasattr(tree, "_fields") else TreeType(subtrees)

so that namedtuples (whose constructor takes one positional argument per field) and tuple subclasses round-trip correctly. tree_merge was still using the plain TreeType(<generator>) form, so merging two namedtuple pytrees binds the whole generator to the first field and raises a misleading TypeError pointing at the merge function:

from mlx.utils import tree_merge
from typing import NamedTuple

class Params(NamedTuple):
    m: int
    b: int

tree_merge(Params(1, 2), Params(10, 20), lambda a, b: a + b)
# TypeError: <lambda>() missing 1 required positional argument: 'b'

This applies the same reconstruction idiom already used by tree_map/tree_map_with_path so namedtuples and tuple subclasses merge correctly, while plain lists and tuples are unchanged.

Changes

• tree_merge reconstructs _fields-bearing containers with TreeType(*subtrees).
• Extended test_supported_trees to cover tree_merge over a namedtuple and a tuple subclass.

Result

tree_merge(Params(1, 2), Params(10, 20), lambda a, b: a + b)
# Params(m=11, b=22)