MLX_UNET.py

[bdeanhardt]

a first pass at the MLX implementation for UNET

[luke-carlson]

example of how to test a pytorch vs mlx implementation, cc: @bdeanhardt @ethanernst11 @levinkhho

Also the config classes, which define the model's structure, are the same so you don't need to copy in any @dataclass configurations, you can import them from the pytorch files as done in the file above.

The benefit there is that the exact same config class can be used for either implementation, which keeps it transparent to the user

[luke-carlson]

Because the test case for MLP_MLX passes, you could start by just merging in that class implantation? If you'd like to merge others, we'd want to make similar test cases comparing them to the pytorch versions before merging