Let me know if youre interested in a collaboration then. I am working on a custom mlx sharding system.