Plans for custom pytrees that work with transformations?

[romanoneg]

Jax has the ability to register custom pytree classes that work with any of their function transformations, are there any plans to add that for mlx or would it cause too much of a headache in edge cases for the moment? It would make custom datastructs possible and enable even more customizability and functionality with mlx.

PS: have been loving mlx so far! thank you :)

[awni]

No plan for that at the moment. One thing you can do as a workaround at the moment is inherit from dict, list or tuple and put what you want to be transformed in inherited container. I think that is sufficiently flexible that it should work for most cases. If you have cases that you think it doesn't cover definitely feel free to share.

[romanoneg]

The main reason for my question is that (and this might be because I just haven't found the right way of doing it) if I create a custom dataclass that inherits from NamedTuple when I vmap over a function that takes in that dataclass it works as it should but I haven't found a way to make it return that custom dataclass out, it instead returns tuples.

While this might not be too much of an issue and casting it after the function is an option, it gets harder to do with custom dataclasses inside of custom dataclasses. Is there a way that I just haven't found to make vmapped-functions output a specific class (unflatten the pytree to a specific custom class)?

[awni]

That's interesting. I hadn't thought of that. Let me take a look and see, maybe there is a simple solution.. maybe not. Not sure yet.