Added support for pytree types that inherit from tuple and typing.namedtuple

[romanoneg]

Proposed changes

I've expanded support for tuple-related datatypes in the transforms and python tree utils.

The reason for this is that there seems to be a hole in the types available that work by default in mlx, at least in how I interact with these types of libraries. While having access to dicts, lists, and tuples is great, if I want to define a custom structure for my data and for it to work natively with the transforms, the closest thing available that has clearly labelled parts is dicts but dicts use more memory overhead, errors with keys can be quieter, and don't have the nice immutability property that namedtuples do.

The changes I've made add this functionality in a very lightweight way, and it doesn't change any previous workflow at all. Additionally, with the changes I also added the ability for types that inherit from tuples to be returned through transformations as it fit perfectly within this new system. Tests for everything were added to demonstrate all these new functionalities.

This helps match the documentation closely as the changes which are in the tree_map and tree_visit_update functions still only focus on default python dict, list and tuple while adding greater support for objects that inherit from tuple now. The docs state:
Functions in this module that return python trees will be using the default python dict, list and tuple but they can usually process objects that inherit from any of these. which I believe fits with the changes made.
If needed, typing.namedtuple can be added to the list above.

As a final benefit, this also helps people migrating from jax as namedtuples are supported by default there.

However, this is my first PR to this repo, if there any glaring mistakes or parts I've overlooked, or if this does not fit within what mlx wants to do, let me know.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[awni]

I think this can be a generator, no need to actually expand a new list here.

[romanoneg]

Good catch, sent the changes.

[awni]

Looks great. I think this is basically good to go. Can you address the comment then I can run the CI?

[awni]

Thanks! It also needs to be rebased to resolve the conflict. If you can do it that'd be great, o/w I'll get to it later.

[romanoneg]

I had some issues moving between different setups with the pre-commit hooks but its all finalized now, however this made me think:
Would it be beneficial if mlx had a way to register custom pytrees? I see namedtuples working with what I want to implement but maybe others would want to create more complex classes?

[awni]

Thanks for the contribution!

Would it be beneficial if mlx had a way to register custom pytrees? I see namedtuples working with what I want to implement but maybe others would want to create more complex classes?

I haven't seen any demand for this.. but if it becomes a common request we can consider it.