Skip to content

MLX_UNET.py #53

Open
bdeanhardt wants to merge 9 commits intoapple:mainfrom
bdeanhardt:bella/mlx
Open

MLX_UNET.py #53
bdeanhardt wants to merge 9 commits intoapple:mainfrom
bdeanhardt:bella/mlx

Conversation

@bdeanhardt
Copy link
Contributor

@bdeanhardt bdeanhardt commented Dec 4, 2024

a first pass at the MLX implementation for UNET

@bdeanhardt bdeanhardt changed the title starting point for MLX_UNET.py MLX_UNET.py Dec 4, 2024
import torch
import numpy as np

def test_pytorch_mlp():
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Copy link
Contributor

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants