I wonder if we could use an @MLXCompile macro like the python decorator to make it easier to write compiled functions? We have some overloads for simple cases (one array in, one array out) but if you don't fall into those cases it is tedious/confusing (I have to look at the tests to figure it out and I am somewhat familiar with this code).

Investigate!