Ressource for adapting 3D maxpool / convTranspose 3D kernels for MLX ?

[arnaudberenbaum]

Hello everyone,
I'm working a lot with 3D medical imaging and I'm very interested in adapting 3D models for image processing for MLX.
From what I understood, MPS does not support 3D convolutions / ConvTranspose / Maxpool3D and would need some new kernel.
Do you have any good ressource if I want to try to create these kernels? :) I've haven't got experience with low level GPU programming but i'm willing to learn 🙂
(just for information, I saw this implementation in pytorch : pytorch/pytorch#116580 )
Thanks a lot !

[awni]

We just landed initial conv3d support in MLX: #993 so you can take a look at that.

[awni]

For Maxpool3D I think you should be able to do it with existing reductions + stride tricks rather than writing a kernel for it.

[arnaudberenbaum]

Ok thanks a lot, I will check this implementation 👍🏼

[francescopisu]

@arnaudberenbaum Any news on this?